diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 045a31f39c..36cfdf9b1e 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -1,11 +1,11 @@ -* @abhinavclemson @gobbleturk @khatwanimohit @bvandermoon @vipannalla @RissyRan @richjames0 @gagika @shralex @SurbhiJainUSC @hengtaoguo @A9isha @aireenmei @NuojCheng @jiangjy1982 @suexu1025 @NicoGrande @jesselu-google @dipannita08 @igorts-git @darisoy @xibinliu +* @abhinavclemson @gobbleturk @khatwanimohit @bvandermoon @vipannalla @RissyRan @richjames0 @gagika @shralex @SurbhiJainUSC @hengtaoguo @A9isha @aireenmei @NuojCheng @jiangjy1982 @suexu1025 @NicoGrande @jesselu-google @dipannita08 @igorts-git @darisoy @xibinliu @huytransformer # Model bring-up -src/maxtext/assets @parambole @shuningjin @RissyRan @suexu1025 @jiangjy1982 @gobbleturk @bvandermoon @gagika @shralex @richjames0 @NicoGrande -src/maxtext/configs/models @parambole @shuningjin @RissyRan @suexu1025 @jiangjy1982 @gobbleturk @bvandermoon @gagika @shralex @richjames0 @NicoGrande @jesselu-google @NuojCheng -src/maxtext/checkpoint_conversion @parambole @shuningjin @RissyRan @suexu1025 @jiangjy1982 @gobbleturk @bvandermoon @hengtaoguo @gagika @shralex @richjames0 @NicoGrande -src/maxtext/layers @parambole @shuningjin @RissyRan @suexu1025 @jiangjy1982 @gobbleturk @bvandermoon @gagika @shralex @richjames0 @NicoGrande @jesselu-google @NuojCheng @xibinliu -src/maxtext/models @parambole @shuningjin @RissyRan @suexu1025 @jiangjy1982 @gobbleturk @bvandermoon @gagika @shralex @richjames0 @NicoGrande @jesselu-google @NuojCheng @xibinliu +src/maxtext/assets @parambole @shuningjin @RissyRan @suexu1025 @jiangjy1982 @gobbleturk @bvandermoon @gagika @shralex @richjames0 @NicoGrande @aireenmei +src/maxtext/configs/models @parambole @shuningjin @RissyRan @suexu1025 @jiangjy1982 @gobbleturk @bvandermoon @gagika @shralex @richjames0 @NicoGrande @jesselu-google @NuojCheng @aireenmei +src/maxtext/checkpoint_conversion @parambole @shuningjin @RissyRan @suexu1025 @jiangjy1982 @gobbleturk @bvandermoon @hengtaoguo @gagika @shralex @richjames0 @NicoGrande @aireenmei +src/maxtext/layers @parambole @shuningjin @RissyRan @suexu1025 @jiangjy1982 @gobbleturk @bvandermoon @gagika @shralex @richjames0 @NicoGrande @jesselu-google @NuojCheng @xibinliu @huytransformer @aireenmei +src/maxtext/models @parambole @shuningjin @RissyRan @suexu1025 @jiangjy1982 @gobbleturk @bvandermoon @gagika @shralex @richjames0 @NicoGrande @jesselu-google @NuojCheng @xibinliu @huytransformer @aireenmei # Features src/maxtext/experimental/rl @A9isha @khatwanimohit @xuefgu @gagika @richjames0 @shralex @NicoGrande @darisoy @@ -29,10 +29,10 @@ src/maxtext/inference @vipannalla @mitalisi @gpolovets1 @mailvijayasingh @jrplat src/dependencies/ @bvandermoon @SurbhiJainUSC @parambole @richjames0 @shralex @darisoy # Docs -docs/ @jacoguzo @bvandermoon @SurbhiJainUSC @richjames0 @shralex @gobbleturk @RissyRan @gagika @A9isha @jiangjy1982 @vipannalla @darisoy +docs/ @jacoguzo @bvandermoon @SurbhiJainUSC @richjames0 @shralex @gobbleturk @RissyRan @gagika @A9isha @jiangjy1982 @vipannalla @darisoy @aireenmei # Workflow files -.github/workflows/ @gobbleturk @khatwanimohit @shralex @parambole @bvandermoon @richjames0 @darisoy @xibinliu +.github/workflows/ @gobbleturk @khatwanimohit @shralex @parambole @bvandermoon @richjames0 @darisoy @xibinliu @huytransformer # Benchmarking/Recipes -benchmarks/ @SujeethJinesh @bvandermoon @richjames0 @shralex @vipannalla @mitalisi @RissyRan @shauryagup @NuojCheng @gobbleturk @khatwanimohit @Obliviour @notabee @suexu1025 @darisoy +benchmarks/ @SujeethJinesh @bvandermoon @richjames0 @shralex @vipannalla @mitalisi @RissyRan @shauryagup @NuojCheng @gobbleturk @khatwanimohit @Obliviour @notabee @suexu1025 @darisoy @aireenmei diff --git a/.github/workflows/build_and_push_docker_image.yml b/.github/workflows/build_and_push_docker_image.yml index b64629d214..15fc192fbe 100644 --- a/.github/workflows/build_and_push_docker_image.yml +++ b/.github/workflows/build_and_push_docker_image.yml @@ -42,6 +42,10 @@ on: required: false type: boolean default: false + run_tests: + required: false + type: boolean + default: true outputs: build_result: description: 'The result of build_and_push job.' @@ -184,7 +188,7 @@ jobs: run_ci_tests: name: Run Unit and Integration Tests needs: [pre_build_check, build_and_push] - if: needs.pre_build_check.outputs.should_run == 'true' && inputs.include_test_assets == true + if: needs.pre_build_check.outputs.should_run == 'true' && inputs.include_test_assets == true && inputs.run_tests == true uses: ./.github/workflows/run_ci_tests.yml with: image_name: ${{ inputs.image_name }} diff --git a/.github/workflows/check_docs_build.yml b/.github/workflows/check_docs_build.yml index a3030ea8ad..5979174805 100644 --- a/.github/workflows/check_docs_build.yml +++ b/.github/workflows/check_docs_build.yml @@ -17,12 +17,11 @@ name: Check the Documentation Build on: - pull_request: - types: [opened, synchronize] workflow_call: inputs: maxtext_sha: - required: true + description: 'The specific MaxText commit SHA.' + required: false type: string workflow_dispatch: @@ -38,7 +37,7 @@ jobs: with: persist-credentials: false fetch-depth: 0 - ref: ${{ inputs.maxtext_sha }} + ref: ${{ inputs.maxtext_sha || github.sha }} - name: Check if only documentation changed id: check diff --git a/.github/workflows/ci_pipeline.yml b/.github/workflows/ci_pipeline.yml index 0f66a51e91..1db52eb13d 100644 --- a/.github/workflows/ci_pipeline.yml +++ b/.github/workflows/ci_pipeline.yml @@ -126,12 +126,23 @@ jobs: echo "run_notebooks=false" >> $GITHUB_OUTPUT fi + code_quality_check: + uses: ./.github/workflows/code_quality.yml + with: + maxtext_sha: ${{ inputs.maxtext_sha || github.sha }} + + docs_build_check: + uses: ./.github/workflows/check_docs_build.yml + with: + maxtext_sha: ${{ inputs.maxtext_sha || github.sha }} + build_and_upload_maxtext_package: - needs: analyze_code_changes - # Run if either tests or notebooks need to run + needs: [analyze_code_changes, code_quality_check, docs_build_check] + # Run if either tests or notebooks need to run; on PRs, gate on code quality + docs passing if: | - needs.analyze_code_changes.outputs.run_tests == 'true' || - needs.analyze_code_changes.outputs.run_notebooks == 'true' + always() && + (needs.analyze_code_changes.outputs.run_tests == 'true' || needs.analyze_code_changes.outputs.run_notebooks == 'true') && + (github.event_name != 'pull_request' || (needs.code_quality_check.result == 'success' && needs.docs_build_check.result == 'success')) uses: ./.github/workflows/build_package.yml with: device_type: tpu @@ -154,15 +165,60 @@ jobs: secrets: HF_TOKEN: ${{ secrets.HF_TOKEN }} + gate_test_run: + name: Gate and Formalize Parameters + needs: [analyze_code_changes, build_and_upload_maxtext_package] + if: | + always() && + needs.analyze_code_changes.outputs.run_tests == 'true' && + needs.build_and_upload_maxtext_package.result == 'success' + runs-on: ubuntu-latest + outputs: + total_workers: ${{ steps.set-params.outputs.total_workers }} + worker_groups: ${{ steps.set-params.outputs.worker_groups }} + steps: + - id: set-params + name: Formalize Test Suite Parameters + run: | + TPU_UNIT_TOTAL_WORKERS=2 + TPU_UNIT_WORKER_GROUPS='[1, 2]' + echo "total_workers=${TPU_UNIT_TOTAL_WORKERS}" >> "$GITHUB_OUTPUT" + echo "worker_groups=${TPU_UNIT_WORKER_GROUPS}" >> "$GITHUB_OUTPUT" + tpu-tests: name: ${{ matrix.flavor || 'TPU' }} tests + needs: [build_and_upload_maxtext_package, gate_test_run] + uses: ./.github/workflows/run_tests_coordinator.yml + strategy: + fail-fast: false + matrix: + flavor: [tpu-unit, tpu-integration, tpu-post-training-unit] + with: + flavor: ${{ matrix.flavor }} + base_image: maxtext-unit-test-tpu:py312 + is_scheduled_run: ${{ github.event_name == 'schedule' }} + maxtext_sha: ${{ needs.build_and_upload_maxtext_package.outputs.maxtext_sha }} + + maxtext_cpu_torch_reference_tests: + name: cpu-torch-reference tests needs: [build_and_upload_maxtext_package] if: needs.analyze_code_changes.outputs.run_tests == 'true' uses: ./.github/workflows/run_tests_coordinator.yml + with: + flavor: cpu-torch-reference + base_image: maxtext-unit-test-tpu:py312 + is_scheduled_run: ${{ github.event_name == 'schedule' }} + maxtext_sha: ${{ needs.build_and_upload_maxtext_package.outputs.maxtext_sha }} + + tpu7x-tests: + name: TPU7X tests + needs: [build_and_upload_maxtext_package, gate_test_run] + if: github.ref == 'refs/heads/main' && (github.event_name == 'schedule' || github.event_name == 'workflow_dispatch') + uses: ./.github/workflows/run_tests_coordinator.yml strategy: fail-fast: false matrix: - flavor: [tpu-unit, tpu-integration, tpu-post-training-unit] + flavor: [tpu7x-unit, tpu7x-integration, tpu7x-post-training-unit] with: flavor: ${{ matrix.flavor }} base_image: maxtext-unit-test-tpu:py312 @@ -171,8 +227,7 @@ jobs: gpu-tests: name: ${{ matrix.flavor || 'GPU' }} tests - needs: [build_and_upload_maxtext_package] - if: needs.analyze_code_changes.outputs.run_tests == 'true' + needs: [build_and_upload_maxtext_package, gate_test_run] strategy: fail-fast: false matrix: @@ -186,8 +241,7 @@ jobs: cpu-tests: name: ${{ matrix.flavor || 'CPU' }} tests - needs: [build_and_upload_maxtext_package] - if: needs.analyze_code_changes.outputs.run_tests == 'true' + needs: [build_and_upload_maxtext_package, gate_test_run] uses: ./.github/workflows/run_tests_coordinator.yml strategy: fail-fast: false @@ -199,31 +253,13 @@ jobs: is_scheduled_run: ${{ github.event_name == 'schedule' }} maxtext_sha: ${{ needs.build_and_upload_maxtext_package.outputs.maxtext_sha }} - setup-pathways-parameters: - name: Setup Pathways Parameters - runs-on: ubuntu-latest - outputs: - total_workers: ${{ steps.set-params.outputs.total_workers }} - worker_groups: ${{ steps.set-params.outputs.worker_groups }} - steps: - - id: set-params - name: Set Pathways Worker Parameters - run: | - # TPU worker constants - TPU_UNIT_TOTAL_WORKERS=2 - TPU_UNIT_WORKER_GROUPS='[1, 2]' - - echo "total_workers=${TPU_UNIT_TOTAL_WORKERS}" >> "$GITHUB_OUTPUT" - echo "worker_groups=${TPU_UNIT_WORKER_GROUPS}" >> "$GITHUB_OUTPUT" - maxtext_tpu_pathways_unit_tests: - needs: [build_and_upload_maxtext_package, setup-pathways-parameters] - if: needs.analyze_code_changes.outputs.run_tests == 'true' + needs: [build_and_upload_maxtext_package, gate_test_run] uses: ./.github/workflows/run_pathways_tests.yml strategy: fail-fast: false matrix: - group: ${{ fromJSON(needs.setup-pathways-parameters.outputs.worker_groups) }} + group: ${{ fromJSON(needs.gate_test_run.outputs.worker_groups || '[1, 2]') }} with: device_type: tpu device_name: v6e-4 @@ -236,11 +272,11 @@ jobs: container_resource_option: "--privileged" is_scheduled_run: ${{ github.event_name == 'schedule' }} maxtext_sha: ${{ needs.build_and_upload_maxtext_package.outputs.maxtext_sha }} - total_workers: ${{ needs.setup-pathways-parameters.outputs.total_workers }} + total_workers: ${{ needs.gate_test_run.outputs.total_workers || '2' }} worker_group: ${{ matrix.group }} maxtext_tpu_pathways_integration_tests: - needs: build_and_upload_maxtext_package + needs: [build_and_upload_maxtext_package, gate_test_run] uses: ./.github/workflows/run_pathways_tests.yml strategy: fail-fast: false @@ -259,20 +295,20 @@ jobs: all_tests_passed: name: All Required Tests Passed - needs: [tpu-tests, gpu-tests, cpu-tests, maxtext_tpu_pathways_unit_tests, maxtext_tpu_pathways_integration_tests] + needs: [build_and_upload_maxtext_package, gate_test_run, tpu-tests, tpu7x-tests, gpu-tests, cpu-tests, maxtext_cpu_torch_reference_tests, maxtext_tpu_pathways_unit_tests, maxtext_tpu_pathways_integration_tests, code_quality_check, docs_build_check] if: always() runs-on: ubuntu-latest steps: - name: Check test results run: | - if [ "${NEEDS_ANALYZE_CHANGES_OUTPUTS_RUN_TESTS}" == "false" ]; then - echo "Tests were skipped" - exit 0 - fi - - # Otherwise, check that build and all tests passed or were skipped + # Check that build and all tests passed or were skipped echo "Build result: ${NEEDS_BUILD_AND_UPLOAD_MAXTEXT_PACKAGE_RESULT}" + echo "Code Quality result: ${NEEDS_CODE_QUALITY_CHECK_RESULT}" + echo "Docs Build result: ${NEEDS_DOCS_BUILD_CHECK_RESULT}" + echo "Gate result: ${NEEDS_GATE_TEST_RUN_RESULT}" echo "TPU Tests (Matrix) result: ${NEEDS_TPU_TESTS_RESULT}" + echo "TPU7X Tests (Matrix) result: ${NEEDS_TPU7X_TESTS_RESULT}" + echo "CPU torch reference tests result: ${NEEDS_MAXTEXT_CPU_TORCH_REFERENCE_TESTS_RESULT}" echo "GPU Tests (Matrix) result: ${NEEDS_GPU_TESTS_RESULT}" echo "CPU Tests (Matrix) result: ${NEEDS_CPU_TESTS_RESULT}" echo "Pathways Unit result: ${NEEDS_MAXTEXT_TPU_PATHWAYS_UNIT_TESTS_RESULT}" @@ -286,10 +322,14 @@ jobs: echo "All required tests passed successfully" env: - NEEDS_ANALYZE_CODE_CHANGES_OUTPUTS_RUN_TESTS: ${{ needs.analyze_code_changes.outputs.run_tests }} NEEDS_BUILD_AND_UPLOAD_MAXTEXT_PACKAGE_RESULT: ${{ needs.build_and_upload_maxtext_package.result }} + NEEDS_CODE_QUALITY_CHECK_RESULT: ${{ needs.code_quality_check.result }} + NEEDS_DOCS_BUILD_CHECK_RESULT: ${{ needs.docs_build_check.result }} + NEEDS_GATE_TEST_RUN_RESULT: ${{ needs.gate_test_run.result }} NEEDS_CPU_TESTS_RESULT: ${{ needs.cpu-tests.result }} NEEDS_TPU_TESTS_RESULT: ${{ needs.tpu-tests.result }} + NEEDS_TPU7X_TESTS_RESULT: ${{ needs.tpu7x-tests.result }} + NEEDS_MAXTEXT_CPU_TORCH_REFERENCE_TESTS_RESULT: ${{ needs.maxtext_cpu_torch_reference_tests.result }} NEEDS_GPU_TESTS_RESULT: ${{ needs.gpu-tests.result }} NEEDS_MAXTEXT_TPU_PATHWAYS_UNIT_TESTS_RESULT: ${{ needs.maxtext_tpu_pathways_unit_tests.result }} NEEDS_MAXTEXT_TPU_PATHWAYS_INTEGRATION_TESTS_RESULT: ${{ needs.maxtext_tpu_pathways_integration_tests.result }} @@ -325,7 +365,7 @@ jobs: notify_failure: name: Notify failed build # creates an issue or modifies last open existing issue for failed build - needs: [tpu-tests, gpu-tests, cpu-tests, maxtext_jupyter_notebooks, maxtext_tpu_pathways_unit_tests, maxtext_tpu_pathways_integration_tests] + needs: [gate_test_run, tpu-tests, tpu7x-tests, gpu-tests, cpu-tests, maxtext_jupyter_notebooks, maxtext_cpu_torch_reference_tests, maxtext_tpu_pathways_unit_tests, maxtext_tpu_pathways_integration_tests, code_quality_check, docs_build_check] if: ${{ always() }} runs-on: ubuntu-latest permissions: @@ -339,7 +379,7 @@ jobs: investigate_failure: name: Investigate failed build # investigates failure of scheduled run and comments on tracking issue - needs: [tpu-tests, gpu-tests, cpu-tests, maxtext_jupyter_notebooks, maxtext_tpu_pathways_unit_tests, maxtext_tpu_pathways_integration_tests, notify_failure] + needs: [gate_test_run, tpu-tests, tpu7x-tests, gpu-tests, cpu-tests, maxtext_jupyter_notebooks, maxtext_cpu_torch_reference_tests, maxtext_tpu_pathways_unit_tests, maxtext_tpu_pathways_integration_tests, code_quality_check, docs_build_check, notify_failure] if: ${{ always() && contains(needs.*.result, 'failure') && github.event_name == 'schedule' }} uses: ./.github/workflows/gemini_investigate.yml permissions: diff --git a/.github/workflows/code_quality.yml b/.github/workflows/code_quality.yml index 6f5c2a277a..187ac62940 100644 --- a/.github/workflows/code_quality.yml +++ b/.github/workflows/code_quality.yml @@ -17,17 +17,12 @@ name: Code Quality on: - pull_request: - -concurrency: - # Dedup pull requests (canceling previous runs of the same workflow for same PR), and scheduled runs but nothing else - group: > - ${{ - github.event_name == 'pull_request' && format('{0}-pr-{1}', github.workflow, github.event.pull_request.number) || - github.event_name == 'schedule' && format('{0}-schedule', github.workflow) || - github.run_id - }} - cancel-in-progress: true + workflow_call: + inputs: + maxtext_sha: + description: 'The specific MaxText commit SHA.' + required: false + type: string jobs: qa: @@ -35,6 +30,9 @@ jobs: runs-on: ubuntu-latest steps: - uses: actions/checkout@v5 + with: + fetch-depth: 0 + ref: ${{ inputs.maxtext_sha || github.sha }} - name: Install uv and set the Python version uses: astral-sh/setup-uv@v7 @@ -56,7 +54,12 @@ jobs: - name: Run pre-commit checks on just the files that have changed run: | - git fetch origin "$GITHUB_BASE_REF":"$GITHUB_BASE_REF" - git branch "$GITHUB_HEAD_REF" . "$GITHUB_WORKSPACE"/venv/bin/activate - pre-commit run --from-ref "$GITHUB_BASE_REF" --to-ref "$GITHUB_HEAD_REF" --show-diff-on-failure + if [ "${{ github.event_name }}" == "pull_request" ]; then + git fetch origin "$GITHUB_BASE_REF":"$GITHUB_BASE_REF" + git branch "$GITHUB_HEAD_REF" + pre-commit run --from-ref "$GITHUB_BASE_REF" --to-ref "$GITHUB_HEAD_REF" --show-diff-on-failure + else + echo "Not a pull request event (${{ github.event_name }}), checking all files." + pre-commit run --all-files --show-diff-on-failure + fi diff --git a/.github/workflows/nightly_images_pipeline.yml b/.github/workflows/nightly_images_pipeline.yml index 4ab0a932b3..32b6a95abe 100644 --- a/.github/workflows/nightly_images_pipeline.yml +++ b/.github/workflows/nightly_images_pipeline.yml @@ -50,7 +50,7 @@ jobs: device_name: v4-8 cloud_runner: linux-x86-n2-16-buildkit - build_and_push_docker_image: + build_and_push_tpu_docker_images: name: Build ${{ matrix.name }} Docker Image needs: build_maxtext_package strategy: @@ -75,6 +75,26 @@ jobs: workflow: post-training image_name: maxtext_post_training_nightly dockerfile: maxtext_tpu_dependencies.Dockerfile + uses: ./.github/workflows/build_and_push_docker_image.yml + with: + image_name: ${{ inputs.image_suffix != '' && format('{0}_{1}', matrix.image_name, inputs.image_suffix) || matrix.image_name }} + device: ${{ matrix.device }} + build_mode: ${{ matrix.build_mode }} + workflow: ${{ matrix.workflow }} + dockerfile: ${{ matrix.dockerfile }} + maxtext_sha: ${{ needs.build_maxtext_package.outputs.maxtext_sha }} + include_test_assets: true + run_tests: false + secrets: + HF_TOKEN: ${{ secrets.HF_TOKEN }} + + build_and_push_gpu_docker_images: + name: Build ${{ matrix.name }} Docker Image + needs: build_maxtext_package + strategy: + fail-fast: false + matrix: + include: - name: "GPU Pre-Training Stable" device: gpu build_mode: stable @@ -94,14 +114,40 @@ jobs: build_mode: ${{ matrix.build_mode }} workflow: ${{ matrix.workflow }} dockerfile: ${{ matrix.dockerfile }} - maxtext_sha: ${{ needs.build_and_upload_maxtext_package.outputs.maxtext_sha }} + maxtext_sha: ${{ needs.build_maxtext_package.outputs.maxtext_sha }} include_test_assets: true secrets: HF_TOKEN: ${{ secrets.HF_TOKEN }} + run_tpu_ci_tests: + name: Run TPU ${{ matrix.name }} CI Tests + needs: build_and_push_tpu_docker_images + strategy: + fail-fast: false + matrix: + include: + - name: "Pre-Training Stable" + image_name: maxtext_jax_stable + device: tpu + workflow: pre-training + - name: "Pre-Training Nightly" + image_name: maxtext_jax_nightly + device: tpu + workflow: pre-training + - name: "Post-Training Nightly" + image_name: maxtext_post_training_nightly + device: tpu + workflow: post-training + uses: ./.github/workflows/run_ci_tests.yml + with: + image_name: ${{ inputs.image_suffix != '' && format('{0}_{1}', matrix.image_name, inputs.image_suffix) || matrix.image_name }} + image_tag: ${{ github.run_id }} + device: ${{ matrix.device }} + workflow: ${{ matrix.workflow }} + run_e2e_tests: name: Run E2E tests - needs: build_and_push_docker_image + needs: build_and_push_tpu_docker_images uses: ./.github/workflows/run_e2e_tests.yml with: mode: ${{ inputs.image_suffix != '' && format('{0}_{1}', 'nightly', inputs.image_suffix) || 'nightly' }} @@ -111,7 +157,7 @@ jobs: notify_failure: name: Notify failed build - needs: [build_and_push_docker_image, run_e2e_tests] + needs: [build_and_push_tpu_docker_images, build_and_push_gpu_docker_images, run_tpu_ci_tests, run_e2e_tests] if: ${{ failure() && inputs.image_suffix == '' }} runs-on: ubuntu-latest permissions: diff --git a/.github/workflows/run_pathways_tests.yml b/.github/workflows/run_pathways_tests.yml index 19adc362c4..921755b060 100644 --- a/.github/workflows/run_pathways_tests.yml +++ b/.github/workflows/run_pathways_tests.yml @@ -113,7 +113,11 @@ jobs: if [ "${{ inputs.is_scheduled_run }}" = "true" ]; then FINAL_PYTEST_MARKER="${{ inputs.pytest_marker }}" else - FINAL_PYTEST_MARKER="${{ inputs.pytest_marker }} and not scheduled_only" + if [ -z "${{ inputs.pytest_marker }}" ]; then + FINAL_PYTEST_MARKER="not scheduled_only or newly_added" + else + FINAL_PYTEST_MARKER="${{ inputs.pytest_marker }} and (not scheduled_only or newly_added)" + fi fi export MAXTEXT_REPO_ROOT=$(pwd) export MAXTEXT_ASSETS_ROOT=$(pwd)/src/maxtext/assets @@ -125,7 +129,7 @@ jobs: PYTHONPATH: "${{ github.workspace }}/src" services: resource_manager: - image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:20260413-jax_0.9.2 + image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:latest ports: - "29001:29001" - "29002:29002" @@ -135,7 +139,7 @@ jobs: TPU_SKIP_MDS_QUERY: true worker: - image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:20260413-jax_0.9.2 + image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/server:latest ports: - "29005:29005" - "29006:29006" @@ -146,7 +150,7 @@ jobs: --tpu=4 proxy: - image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server:20260413-jax_0.9.2 + image: us-docker.pkg.dev/cloud-tpu-v2-images/pathways/proxy_server:latest ports: - "29000:29000" env: diff --git a/.github/workflows/run_tests_against_package.yml b/.github/workflows/run_tests_against_package.yml index 4b82a118e0..5ae344cd76 100644 --- a/.github/workflows/run_tests_against_package.yml +++ b/.github/workflows/run_tests_against_package.yml @@ -74,6 +74,11 @@ on: required: false type: boolean default: false + install_torch_cpu: + description: 'Install CPU-only torch for layer-wise reference tests' + required: false + type: boolean + default: false permissions: contents: read jobs: @@ -100,6 +105,7 @@ jobs: uses: actions/checkout@v5 with: ref: ${{ inputs.maxtext_sha }} + fetch-depth: 0 - name: Download the maxtext wheel if: ${{ !inputs.maxtext_installed }} uses: actions/download-artifact@634f93cb2916e3fdff6788551b99b062d0335ce0 # v5.0.0 @@ -119,8 +125,14 @@ jobs: else install_tpu_pre_train_extra_deps fi + if [ "${INPUTS_INSTALL_TORCH_CPU}" == "true" ]; then + python3 -m pip install --quiet torch==2.11.0 --index-url https://download.pytorch.org/whl/cpu --no-deps + python3 -c "import torch; print(f'torch {torch.__version__}')" + fi python3 --version python3 -m pip freeze + env: + INPUTS_INSTALL_TORCH_CPU: ${{ inputs.install_torch_cpu }} - name: Copy test assets files if: ${{ !inputs.maxtext_installed }} run : gcloud storage cp gs://maxtext-test-assets/* tests/assets @@ -147,7 +159,11 @@ jobs: if [ "${INPUTS_IS_SCHEDULED_RUN}" == "true" ]; then FINAL_PYTEST_MARKER="${INPUTS_PYTEST_MARKER}" else - FINAL_PYTEST_MARKER="${INPUTS_PYTEST_MARKER} and not scheduled_only" + if [ -z "${INPUTS_PYTEST_MARKER}" ]; then + FINAL_PYTEST_MARKER="not scheduled_only or newly_added" + else + FINAL_PYTEST_MARKER="${INPUTS_PYTEST_MARKER} and (not scheduled_only or newly_added)" + fi fi # TODO: Use package data for testing and remove the env vars export MAXTEXT_REPO_ROOT=$(pwd) diff --git a/.github/workflows/run_tests_coordinator.yml b/.github/workflows/run_tests_coordinator.yml index 1911d8d98f..a47ae59655 100644 --- a/.github/workflows/run_tests_coordinator.yml +++ b/.github/workflows/run_tests_coordinator.yml @@ -25,9 +25,12 @@ on: Test flavor ( tpu-unit, tpu-integration, tpu-post-training-unit, tpu-post-training-integration, + tpu7x-unit, tpu7x-integration, + tpu7x-post-training-unit, tpu7x-post-training-integration, gpu-unit, gpu-integration, cpu-unit, - cpu-post-training-unit + cpu-post-training-unit, + cpu-torch-reference ) required: true type: string @@ -113,10 +116,15 @@ jobs: "tpu-integration": "tpu", "tpu-post-training-unit": "tpu", "tpu-post-training-integration": "tpu", + "tpu7x-unit": "tpu", + "tpu7x-integration": "tpu", + "tpu7x-post-training-unit": "tpu", + "tpu7x-post-training-integration": "tpu", "gpu-unit": "cuda12", "gpu-integration": "cuda12", "cpu-unit": "cpu", - "cpu-post-training-unit": "cpu" + "cpu-post-training-unit": "cpu", + "cpu-torch-reference": "cpu" }')[inputs.flavor] }} device_name: >- @@ -125,10 +133,15 @@ jobs: "tpu-integration": "v6e-4", "tpu-post-training-unit": "v6e-4", "tpu-post-training-integration": "v6e-4", + "tpu7x-unit": "tpu7x-8", + "tpu7x-integration": "tpu7x-8", + "tpu7x-post-training-unit": "tpu7x-8", + "tpu7x-post-training-integration": "tpu7x-8", "gpu-unit": "a100-40gb-4", "gpu-integration": "a100-40gb-4", "cpu-unit": "X64", - "cpu-post-training-unit": "X64" + "cpu-post-training-unit": "X64", + "cpu-torch-reference": "X64" }')[inputs.flavor] }} cloud_runner: >- @@ -137,10 +150,15 @@ jobs: "tpu-integration": "linux-x86-ct6e-180-4tpu", "tpu-post-training-unit": "linux-x86-ct6e-180-4tpu", "tpu-post-training-integration": "linux-x86-ct6e-180-4tpu", + "tpu7x-unit": "linux-x86-tpu7x-224-4tpu", + "tpu7x-integration": "linux-x86-tpu7x-224-4tpu", + "tpu7x-post-training-unit": "linux-x86-tpu7x-224-4tpu", + "tpu7x-post-training-integration": "linux-x86-tpu7x-224-4tpu", "gpu-unit": "linux-x86-a2-48-a100-4gpu", "gpu-integration": "linux-x86-a2-48-a100-4gpu", "cpu-unit": "linux-x86-n2-32", - "cpu-post-training-unit": "linux-x86-n2-32" + "cpu-post-training-unit": "linux-x86-n2-32", + "cpu-torch-reference": "linux-x86-n2-32" }')[inputs.flavor] }} # Pytest Marker Mapping pytest_marker: >- @@ -149,10 +167,15 @@ jobs: "tpu-integration": "not cpu_only and not gpu_only and integration_test and not post_training", "tpu-post-training-unit": "not cpu_only and not gpu_only and not integration_test and post_training", "tpu-post-training-integration": "not cpu_only and not gpu_only and integration_test", + "tpu7x-unit": "not cpu_only and not gpu_only and not integration_test and not post_training", + "tpu7x-integration": "not cpu_only and not gpu_only and integration_test and not post_training", + "tpu7x-post-training-unit": "not cpu_only and not gpu_only and not integration_test and post_training", + "tpu7x-post-training-integration": "not cpu_only and not gpu_only and integration_test", "gpu-unit": "not cpu_only and not tpu_only and not integration_test and not post_training", "gpu-integration": "not cpu_only and not tpu_only and integration_test and not post_training", "cpu-unit": "cpu_only and not post_training", - "cpu-post-training-unit": "cpu_only and post_training" + "cpu-post-training-unit": "cpu_only and post_training", + "cpu-torch-reference": "not post_training" }')[inputs.flavor] }} pytest_addopts: >- @@ -161,10 +184,15 @@ jobs: "tpu-integration": "", "tpu-post-training-unit": "tests/post_training/unit tests/unit", "tpu-post-training-integration": "tests/post_training/integration", + "tpu7x-unit": "", + "tpu7x-integration": "", + "tpu7x-post-training-unit": "tests/post_training/unit tests/unit", + "tpu7x-post-training-integration": "tests/post_training/integration", "gpu-unit": "", "gpu-integration": "", "cpu-unit": "", - "cpu-post-training-unit": "tests/post_training/unit tests/unit" + "cpu-post-training-unit": "tests/post_training/unit tests/unit", + "cpu-torch-reference": "" }')[inputs.flavor] }} pytest_extra_args: >- @@ -173,10 +201,15 @@ jobs: "tpu-integration": "--ignore=tests/post_training", "tpu-post-training-unit": "", "tpu-post-training-integration": "", + "tpu7x-unit": "--ignore=tests/post_training", + "tpu7x-integration": "--ignore=tests/post_training", + "tpu7x-post-training-unit": "", + "tpu7x-post-training-integration": "", "gpu-unit": "--ignore=tests/post_training", "gpu-integration": "--ignore=tests/post_training", "cpu-unit": "--ignore=tests/post_training", - "cpu-post-training-unit": "" + "cpu-post-training-unit": "", + "cpu-torch-reference": "-o addopts= -rf --import-mode=importlib --strict-markers tests/unit/gemma4_layers_test.py tests/unit/gemma4_small_layers_test.py tests/unit/qwen3_next_vs_reference_test.py tests/unit/qwen3_5_layers_test.py tests/unit/qwen3_omni_layers_test.py" }')[inputs.flavor] }} ${{ inputs.additional_pytest_args }} @@ -197,3 +230,4 @@ jobs: total_workers: ${{ fromJSON(needs.setup-parameters.outputs.total_workers) }} maxtext_sha: ${{ inputs.maxtext_sha }} is_update_hlo: ${{ inputs.is_update_hlo }} + install_torch_cpu: ${{ inputs.flavor == 'cpu-torch-reference' }} diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index e43a765023..823c48103d 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -31,9 +31,6 @@ repos: - '--disable=R0401,R0917,W0201,W0613' - "--ignore-patterns='.pytype,.*pyi$'" - '--ignore-paths=src/maxtext/input_pipeline/protos' - - 'benchmarks' - - 'src' - - 'tests' # - repo: https://github.com/google/pytype # rev: 2024.10.11 diff --git a/docs/development/update_dependencies.md b/docs/development/update_dependencies.md index f1f4fa04e0..e921b0b268 100644 --- a/docs/development/update_dependencies.md +++ b/docs/development/update_dependencies.md @@ -80,7 +80,7 @@ overrides. You will need to do this separately for the TPU and GPU environments. ### TPU Pre-Training -> **Note:** The current `src/dependencies/requirements/generated_requirements/tpu-requirements.txt` in the repository was generated using JAX build commit hash: [efd6cf797ee9c4f29c6c6d5e91ae4432209063be](https://github.com/jax-ml/jax/commit/efd6cf797ee9c4f29c6c6d5e91ae4432209063be). When regenerating the requirements, either use the same commit hash or update this hash if you use a different one. +> **Note:** The current `src/dependencies/requirements/generated_requirements/tpu-requirements.txt` in the repository was generated using JAX build commit hash: [31c559c7744105009f5d5fcfc3919e505db37f8b](https://github.com/jax-ml/jax/commit/31c559c7744105009f5d5fcfc3919e505db37f8b). When regenerating the requirements, either use the same commit hash or update this hash if you use a different one. If you have made changes to TPU pre-training dependencies in `src/dependencies/requirements/base_requirements/tpu-requirements.txt`, you need to regenerate the pinned pre-training requirements in `generated_requirements/` directory. Run the following command, replacing `` with the hash you copied in the previous step: @@ -89,7 +89,7 @@ bash src/dependencies/scripts/generate_requirements.sh \ --base-requirements src/dependencies/requirements/base_requirements/tpu-requirements.txt \ --generated-requirements tpu-requirements.txt \ --override-requirements src/dependencies/extra_deps/tpu_overrides.txt \ ---seed-commit efd6cf797ee9c4f29c6c6d5e91ae4432209063be +--seed-commit 31c559c7744105009f5d5fcfc3919e505db37f8b # Copy generated requirements to src/dependencies/requirements/generated_requirements mv generated_artifacts/python3_12/tpu-requirements.txt \ diff --git a/docs/guides/checkpointing_solutions/convert_checkpoint.md b/docs/guides/checkpointing_solutions/convert_checkpoint.md index e65627e794..4f281d5b24 100644 --- a/docs/guides/checkpointing_solutions/convert_checkpoint.md +++ b/docs/guides/checkpointing_solutions/convert_checkpoint.md @@ -78,7 +78,7 @@ You can find your converted checkpoint files under `${BASE_OUTPUT_DIRECTORY}/0/i ### Key Parameters - `model_name`: The specific model identifier. It must match a supported entry in the MaxText [globals.py](https://github.com/AI-Hypercomputer/maxtext/blob/16b684840db9b96b19e24e84ac49f06af7204ae3/src/maxtext/utils/globals.py#L46C1-L46C7). -- `scan_layers`: Controls whether the output uses a scanned (`scan_layers=true`) or unscanned (`scan_layers=false`) checkpoint format. Refer [to the Checkpoints guide](checkpoints) for more information. **IMPORTANT:** This setting *must* match the `scan_layers` value used during model training or loading. A mismatch will cause PyTree loading errors (though MaxText will intercept these and raise a descriptive `ValueError` explaining the mismatch). +- `scan_layers`: Controls whether the output uses a scanned (`scan_layers=true`) or unscanned (`scan_layers=false`) checkpoint format. Refer [to the Checkpoints guide](checkpoints) for more information. **Note:** When resuming or loading a checkpoint in MaxText, this setting will automatically be loaded from the checkpoint metadata, meaning you do not have to manually specify it during model execution. If you do explicitly specify a value for `scan_layers`, it must match the checkpoint's saved configuration, or a `ValueError` mismatch error will be raised. - `use_multimodal`: Indicates if multimodality is used, important for Gemma3. - `base_output_directory`: The path where the converted Orbax checkpoint will be stored; it can be Google Cloud Storage (GCS) or local. - `hardware=cpu`: The conversion script runs on a CPU machine. @@ -298,9 +298,9 @@ Here is an example [PR to add support for gemma3 multi-modal model](https://gith - **Error:** `Type ShapeDtypeStruct is not a valid JAX type` or generic **PyTree structure/shape mismatches** (e.g., Orbax reporting `"X/Y paths matched"`, such as `143/145 paths`). - - **Cause: Configuration mismatch** (e.g., `scan_layers`) between the checkpoint conversion script (e.g., `to_maxtext.py` or `to_huggingface.py`) and the trainer/inference runner (e.g., `train.py`). + - **Cause: Configuration mismatch** (e.g., `scan_layers`) between the checkpoint conversion script (e.g., `to_maxtext.py` or `to_huggingface.py`) and the trainer/inference runner (e.g., `train.py`). (Note: Since MaxText automatically loads `scan_layers` from the checkpoint's saved metadata, you should only encounter this error if you explicitly set a mismatching value on the command line, which raises a `ValueError` mismatch error). - - **Solution:** Ensure the `scan_layers` flag is set to the exact same value (`True` or `False`) in both the conversion command and your training/execution command. + - **Solution:** Omit the `scan_layers` parameter from your training or execution command to allow MaxText to automatically resolve it from the checkpoint metadata, or ensure any explicitly specified `scan_layers` parameter matches the format of the loaded checkpoint. - **Error:** The converted checkpoint loads without errors but produces nonsensical output. diff --git a/docs/reference/core_concepts/checkpoints.md b/docs/reference/core_concepts/checkpoints.md index c365572688..899a924b90 100644 --- a/docs/reference/core_concepts/checkpoints.md +++ b/docs/reference/core_concepts/checkpoints.md @@ -74,7 +74,7 @@ In MaxText, the **`scan_layers`** configuration parameter is used to control thi - `scan_layers=false` tells MaxText to keep layer parameters unstacked (often required for inference and certain model architectures). > [!IMPORTANT] -> **PyTree Structure Compatibility:** Because JAX expects the loaded PyTree structure to exactly match the model's instantiated structure, the value of the `scan_layers` flag during execution (training, SFT, RL, DPO, or decoding) **must** match the format of the checkpoint being loaded. A mismatch will cause PyTree loading or shape/path mismatch errors (which MaxText will intercept to raise a descriptive `ValueError` pointing to the scan_layers setting). +> **Automatic scan_layers Resolution:** MaxText automatically loads `scan_layers` from the checkpoint's saved metadata when resuming (via `load_parameters_path`) if you do not explicitly specify `scan_layers` on the command-line. If you explicitly specify a value for `scan_layers` that conflicts with the checkpoint format, MaxText will raise a descriptive `ValueError` mismatch error to prevent JAX PyTree structure or shape mismatch errors during loading. ### Takeaways diff --git a/docs/run_maxtext/run_maxtext_localhost.md b/docs/run_maxtext/run_maxtext_localhost.md index 0689ac956c..3f1ba692c2 100644 --- a/docs/run_maxtext/run_maxtext_localhost.md +++ b/docs/run_maxtext/run_maxtext_localhost.md @@ -70,7 +70,7 @@ python3 -m maxtext.inference.decode \ **Note:** Because the model hasn't been properly trained, the output text will be random. To generate meaningful output, you need to load a trained checkpoint using the `load_parameters_path` argument. For instructions on how to convert pre-trained Hugging Face model checkpoints (like Llama or Gemma) to MaxText's Orbax format, please refer to the [Checkpoint Conversion Guide](../guides/checkpointing_solutions/convert_checkpoint.md). > [!NOTE] -> **Checkpoints & `scan_layers` compatibility:** When loading an external or converted checkpoint via `load_parameters_path`, the `scan_layers` setting in your command **must** match the setting used to save the checkpoint. If the checkpoint was saved/converted with `scan_layers=False` (common for Hugging Face conversions and inference runs), you must specify `scan_layers=False` in your command. Otherwise, JAX/Orbax will raise PyTree structure mismatch errors. +> **Checkpoints & `scan_layers` Auto-Resolution:** When resuming or loading an external or converted checkpoint via `load_parameters_path`, MaxText automatically loads `scan_layers` with the checkpoint's saved format (stacked vs unstacked). You do not need to explicitly specify `scan_layers` on your command line when resuming. If you do explicitly specify a `scan_layers` value, it must match the checkpoint's saved configuration, or a `ValueError` mismatch error will be raised. ### Running models using provided configs diff --git a/docs/tutorials/posttraining/dpo.md b/docs/tutorials/posttraining/dpo.md index 5593e96742..039b6c9b25 100644 --- a/docs/tutorials/posttraining/dpo.md +++ b/docs/tutorials/posttraining/dpo.md @@ -104,11 +104,11 @@ export MAXTEXT_CKPT_PATH= # e.g., gs://my-bucket/my-model-checkpoint/ ``` > [!IMPORTANT] -> **Matching the `scan_layers` Parameter:** -> The `scan_layers` setting during your fine-tuning run **must match** the setting used when creating the checkpoint at `MAXTEXT_CKPT_PATH`. +> **Automatic `scan_layers` Resolution:** +> MaxText automatically loads `scan_layers` from the checkpoint's saved metadata when resuming (via `load_parameters_path`) if you do not explicitly specify it on the command-line. > -> - If the checkpoint was converted or saved with `scan_layers=False` (which is common for Hugging Face conversions and inference-ready models), you **must also provide `scan_layers=False` in the MaxText command.** -> - If `scan_layers` does not match, MaxText will raise a `ValueError`. +> - You do not need to manually supply `scan_layers=False` (or `scan_layers=True`) when loading checkpoints; MaxText will configure this automatically. +> - If you do explicitly provide a `scan_layers` argument, it must match the checkpoint's saved setting or a `ValueError` mismatch error will be raised. > See the [Checkpoints concept guide](../../reference/core_concepts/checkpoints.md) for more details. ## Running DPO Training diff --git a/docs/tutorials/posttraining/multimodal.md b/docs/tutorials/posttraining/multimodal.md index 0e3fe53a04..90ceff00bc 100644 --- a/docs/tutorials/posttraining/multimodal.md +++ b/docs/tutorials/posttraining/multimodal.md @@ -101,10 +101,11 @@ export MAXTEXT_CKPT_PATH= # e.g., gs://my-bucket/my-model-checkpoint/ ``` > [!IMPORTANT] -> **Matching the `scan_layers` Parameter:** -> The `scan_layers` setting during your fine-tuning run **must match** the setting used when creating the checkpoint at `MAXTEXT_CKPT_PATH`. +> **Automatic `scan_layers` Resolution:** +> MaxText automatically loads `scan_layers` from the checkpoint's saved metadata when resuming (via `load_parameters_path`) if you do not explicitly specify it on the command-line. > -> - If the checkpoint was converted or saved with `scan_layers=False` (which is common for Hugging Face conversions and inference-ready models), you **must also provide `scan_layers=False` in the MaxText command.** +> - You do not need to manually supply `scan_layers=False` (or `scan_layers=True`) when loading checkpoints; MaxText will configure this automatically. +> - If you do explicitly provide a `scan_layers` argument, it must match the checkpoint's saved setting or a `ValueError` mismatch error will be raised. ## Multimodal Decode diff --git a/docs/tutorials/posttraining/rl_on_multi_host.md b/docs/tutorials/posttraining/rl_on_multi_host.md index bc53bcf559..9cca5bf3fa 100644 --- a/docs/tutorials/posttraining/rl_on_multi_host.md +++ b/docs/tutorials/posttraining/rl_on_multi_host.md @@ -149,11 +149,11 @@ export MAXTEXT_CKPT_PATH= # e.g., gs://my-bucket/my-model-checkpoint/ ``` > [!IMPORTANT] -> **Matching the `scan_layers` Parameter:** -> The `scan_layers` setting during your RL training run **must match** the setting used when creating the checkpoint at `MAXTEXT_CKPT_PATH`. +> **Automatic `scan_layers` Resolution:** +> MaxText automatically loads `scan_layers` from the checkpoint's saved metadata when resuming (via `load_parameters_path`) if you do not explicitly specify it on the command-line. > -> - If the checkpoint was converted or saved with `scan_layers=False` (which is common for Hugging Face conversions and inference-ready models), you **must also provide `scan_layers=False` in the MaxText command.** -> - If `scan_layers` does not match, MaxText will raise a `ValueError`. +> - You do not need to manually supply `scan_layers=False` (or `scan_layers=True`) when loading checkpoints; MaxText will configure this automatically. +> - If you do explicitly provide a `scan_layers` argument, it must match the checkpoint's saved setting or a `ValueError` mismatch error will be raised. > See the [Checkpoints concept guide](../../reference/core_concepts/checkpoints.md) for more details. ## Submit your RL workload via Pathways diff --git a/docs/tutorials/posttraining/sft.md b/docs/tutorials/posttraining/sft.md index eba5a8f76c..61bf68af84 100644 --- a/docs/tutorials/posttraining/sft.md +++ b/docs/tutorials/posttraining/sft.md @@ -89,11 +89,11 @@ export MAXTEXT_CKPT_PATH= # e.g., gs://my-bucket/my-model-checkpoint/ ``` > [!IMPORTANT] -> **Matching the `scan_layers` Parameter:** -> The `scan_layers` setting during your fine-tuning run **must match** the setting used when creating the checkpoint at `MAXTEXT_CKPT_PATH`. +> **Automatic `scan_layers` Resolution:** +> MaxText automatically loads `scan_layers` from the checkpoint's saved metadata when resuming (via `load_parameters_path`) if you do not explicitly specify it on the command-line. > -> - If the checkpoint was converted or saved with `scan_layers=False` (which is common for Hugging Face conversions and inference-ready models), you **must also provide `scan_layers=False` in the MaxText command.** -> - If `scan_layers` does not match, MaxText will raise a `ValueError`. +> - You do not need to manually supply `scan_layers=False` (or `scan_layers=True`) when loading checkpoints; MaxText will configure this automatically. +> - If you do explicitly provide a `scan_layers` argument, it must match the checkpoint's saved setting or a `ValueError` mismatch error will be raised. > See the [Checkpoints concept guide](../../reference/core_concepts/checkpoints.md) for more details. ## Run SFT on Hugging Face Dataset diff --git a/docs/tutorials/posttraining/sft_on_multi_host.md b/docs/tutorials/posttraining/sft_on_multi_host.md index 4b7805d4b1..db0d876e5a 100644 --- a/docs/tutorials/posttraining/sft_on_multi_host.md +++ b/docs/tutorials/posttraining/sft_on_multi_host.md @@ -140,11 +140,11 @@ export MAXTEXT_CKPT_PATH= # gs://my-bucket/my-checkpoint-directory/0/ ``` > [!IMPORTANT] -> **Matching the `scan_layers` Parameter:** -> The `scan_layers` setting during your fine-tuning run **must match** the setting used when creating the checkpoint at `MAXTEXT_CKPT_PATH`. +> **Automatic `scan_layers` Resolution:** +> MaxText automatically loads `scan_layers` from the checkpoint's saved metadata when resuming (via `load_parameters_path`) if you do not explicitly specify it on the command-line. > -> - If the checkpoint was converted or saved with `scan_layers=False` (which is common for Hugging Face conversions and inference-ready models), you **must also provide `scan_layers=False` in the MaxText command.** -> - If `scan_layers` does not match, MaxText will raise a `ValueError`. +> - You do not need to manually supply `scan_layers=False` (or `scan_layers=True`) when loading checkpoints; MaxText will configure this automatically. +> - If you do explicitly provide a `scan_layers` argument, it must match the checkpoint's saved setting or a `ValueError` mismatch error will be raised. > See the [Checkpoints concept guide](../../reference/core_concepts/checkpoints.md) for more details. ## Submit workload on GKE cluster diff --git a/pytest.ini b/pytest.ini index 4976e05ebc..9cfe0e6103 100644 --- a/pytest.ini +++ b/pytest.ini @@ -38,6 +38,7 @@ markers = NOTE: this marker is not to be used manually, it is auto- applied to tests with external_* or tpu_only marker. scheduled_only: marks tests to run only on scheduled CI runs + newly_added: newly introduced or modified tests in PRs, executed even if scheduled_only integration_test: tests exercising larger portions of the system, including interactions with other systems like GCS, e.g., end_to_end tests diff --git a/scripts/diloco/run_diloco.sh b/scripts/diloco/run_diloco.sh new file mode 100644 index 0000000000..763a01ccd4 --- /dev/null +++ b/scripts/diloco/run_diloco.sh @@ -0,0 +1,101 @@ +#!/bin/bash + +# This script launches a DiLoCo pre-training workload on a GKE cluster using XPK. + +set -e + +# --- Cluster Parameters --- +export PROJECT_ID="${PROJECT_ID:-}" +export CLUSTER_NAME="${CLUSTER_NAME:-}" +export ZONE="${ZONE:-}" +export RESERVATION="${RESERVATION:-}" # optional +export BASE_OUTPUT_DIRECTORY="${BASE_OUTPUT_DIRECTORY:-}" # change to your own GCS bucket for logging and checkpointing +export DATASET_PATH="${DATASET_PATH:-}" # change to your own GSC bucket for datasets. Make sure datasets exists +export DOCKER_IMAGE="${DOCKER_IMAGE:-}" # Full path to the Docker image you pushed (e.g., gcr.io/tpu-prod-env-multipod/maxtext_jax_stable:2026-06-22) +export TPU_TYPE="${TPU_TYPE:-}" # At least v5p-32 is needed to run Qwen3-30b-a3b. v5p-8 for qwen3-8b +export NUM_SLICES="${NUM_SLICES:-}" # you need at least two slices to let diloco take effect +export WORKLOAD_NAME="${WORKLOAD_NAME:-$(whoami)-diloco-${TPU_TYPE}-$(date +%Y%m%d-%H%M%S)}" # this will be the name of run, for logging purposes + +# --- Model Parameters --- +export MODEL_NAME="${MODEL_NAME:-}" +export PER_DEVICE_BATCH_SIZE="${PER_DEVICE_BATCH_SIZE:-8}" +export MAX_TARGET_LENGTH="${MAX_TARGET_LENGTH:-2048}" +export TRAINING_STEPS="${TRAINING_STEPS:-20}" + +# --- DiLoCo Parameters --- +export DILOCO_SYNC_PERIOD="${DILOCO_SYNC_PERIOD:-10}" +export DILOCO_OUTER_LR="${DILOCO_OUTER_LR:-0.1}" +export DILOCO_OUTER_MOMENTUM="${DILOCO_OUTER_MOMENTUM:-0.9}" + +# --- XLA Flags --- +export XLA_FLAGS="${XLA_FLAGS:- \ + --xla_tpu_scoped_vmem_limit_kib=65536 \ +--xla_tpu_bf16_emission_mode=NATIVE_EMISSION \ +--xla_tpu_enable_sparse_core_reduce_scatter_v2=true \ +--xla_tpu_enable_sparse_core_collective_offload_all_gather=true \ +--xla_tpu_enable_sparse_core_collective_offload_2d_all_gather=true \ +--xla_tpu_enable_all_gather_offload_tracing=true \ +--xla_tpu_use_tc_device_shape_on_sc=True \ +--xla_sc_disable_megacore_partitioning=True \ +--xla_tpu_enable_async_collective_fusion_fuse_all_gather=false \ +--xla_enable_async_all_gather=true \ +--xla_tpu_prefer_async_allgather_to_allreduce=true \ +--xla_tpu_enable_sparse_core_collective_offload_all_reduce=true \ +--xla_tpu_enable_sparse_core_collective_offload_reduce_scatter=true \ +--xla_tpu_enable_sparse_core_collective_offload_3d_all_gather=true \ +--xla_tpu_use_single_sparse_core_for_all_gather_offload=true \ +--xla_tpu_enable_concurrent_sparse_core_offloading=true \ +--xla_tpu_aggressive_opt_barrier_removal=true \ +--xla_tpu_enable_offloading_gather_to_sparsecore=true \ +--xla_tpu_sparse_core_all_gather_latency_multiplier=1 \ +--xla_tpu_sparse_core_reduce_scatter_latency_multiplier=3 \ +--xla_tpu_enable_sparse_core_collective_aggregator=true \ +--xla_tpu_enable_latency_hiding_layer_scheduler=true \ +--xla_tpu_scheduler_percent_shared_memory_limit=150 \ +--xla_tpu_enable_layer_scheduler_for_dependent_collectives=true \ +--xla_tpu_enable_sparse_core_collective_offload_nd_reduce_scatter=true \ +--xla_tpu_pcie_bandwidth_multiplier=0.03 \ +--xla_tpu_enable_sparse_core_offload_queuing_in_lhs=true \ +--xla_tpu_enable_multi_compute_overlap_in_layer_scheduler=false \ + --xla_tpu_enable_3d_reduce_scatter_decomposer=false }" + +if [ "$NUM_SLICES" -lt 2 ]; then + echo "Warning: NUM_SLICES is less than 2. DiLoCo will not take effect." +fi + +# MaxText command +MAXTEXT_COMMAND="cd /deps/src/ && \ +LIBTPU_INIT_ARGS='${XLA_FLAGS}' \ +python3 -m maxtext.trainers.pre_train.train \ +maxtext/configs/base.yml \ +run_name=$WORKLOAD_NAME \ +save_config_to_gcs=true \ +base_output_directory=$BASE_OUTPUT_DIRECTORY \ +dataset_path=$DATASET_PATH \ +dataset_name='c4/en:3.0.1' \ +eval_dataset_name='c4/en:3.0.1' \ +model_name=$MODEL_NAME \ +tokenizer_type=huggingface \ +tokenizer_path=maxtext/assets/tokenizers/qwen3-tokenizer \ +per_device_batch_size=$PER_DEVICE_BATCH_SIZE \ +max_target_length=$MAX_TARGET_LENGTH \ +enable_diloco=true \ +dcn_diloco_parallelism=$NUM_SLICES \ +diloco_sync_period=$DILOCO_SYNC_PERIOD \ +diloco_outer_lr=$DILOCO_OUTER_LR \ +diloco_outer_momentum=$DILOCO_OUTER_MOMENTUM \ +steps=$TRAINING_STEPS" + +# Workload Creation +xpk workload create \ + --cluster="$CLUSTER_NAME" \ + --project="$PROJECT_ID" \ + --zone="$ZONE" \ + --priority=medium \ + --max-restarts=0 \ + --tpu-type="$TPU_TYPE" \ + --num-slices="$NUM_SLICES" \ + --docker-image="${DOCKER_IMAGE}" \ + --workload="${WORKLOAD_NAME}" \ + ${RESERVATION:+--reservation="$RESERVATION"} \ + --command="${MAXTEXT_COMMAND}" diff --git a/src/dependencies/requirements/base_requirements/requirements.txt b/src/dependencies/requirements/base_requirements/requirements.txt index 5ba8ee5093..68a4384e97 100644 --- a/src/dependencies/requirements/base_requirements/requirements.txt +++ b/src/dependencies/requirements/base_requirements/requirements.txt @@ -46,3 +46,8 @@ tokamax!=0.1.0 transformers>=5.8.0 uvloop qwix>=0.1.6 +coverage +blobfile +shortuuid +nltk +evaluate diff --git a/src/dependencies/requirements/generated_requirements/cuda12-requirements.txt b/src/dependencies/requirements/generated_requirements/cuda12-requirements.txt index 4e52d9125a..266e7b3ff9 100644 --- a/src/dependencies/requirements/generated_requirements/cuda12-requirements.txt +++ b/src/dependencies/requirements/generated_requirements/cuda12-requirements.txt @@ -3,13 +3,13 @@ absl-py>=2.4.0 aiofiles>=25.1.0 -aiohappyeyeballs>=2.6.2 -aiohttp>=3.13.5 +aiohappyeyeballs>=2.7.1 +aiohttp>=3.14.1 aiosignal>=1.4.0 annotated-doc>=0.0.4 annotated-types>=0.7.0 antlr4-python3-runtime>=4.9.3 -anyio>=4.13.0 +anyio>=4.14.1 aqtp>=0.9.0 array-record>=0.8.3 astroid>=4.0.4 @@ -17,34 +17,37 @@ astunparse>=1.6.3 attrs>=26.1.0 auditwheel>=6.6.0 black>=25.12.0 +blobfile>=3.2.0 build>=1.4.3 certifi>=2026.2.25 cffi>=2.0.0 ; platform_python_implementation != 'PyPy' cfgv>=3.5.0 charset-normalizer>=3.4.7 -chex>=0.1.91 -click>=8.4.0 +chex>=0.1.92 +click>=8.4.2 cloud-accelerator-diagnostics>=0.1.1 cloudpickle>=3.1.2 clu>=0.0.12 colorama>=0.4.6 contourpy>=1.3.3 -cryptography>=48.0.0 +coverage>=7.14.3 +cryptography>=49.0.0 cycler>=0.12.1 -datasets>=4.8.5 +datasets>=5.0.0 decorator>=5.3.1 dill>=0.4.1 -distlib>=0.4.0 +distlib>=0.4.3 distro>=1.9.0 dm-tree>=0.1.10 docstring-parser>=0.18.0 -drjax>=0.1.4 +drjax>=0.2.0 editdistance>=0.8.1 einops>=0.8.2 einshape>=1.0 etils>=1.14.0 +evaluate>=0.4.6 execnet>=2.1.2 -fastapi>=0.136.1 +fastapi>=0.139.0 filelock>=3.28.0 flatbuffers>=25.12.19 flax>=0.12.7 @@ -53,40 +56,40 @@ frozenlist>=1.8.0 fsspec>=2026.2.0 gast>=0.7.0 gcsfs>=2026.2.0 -google-api-core>=2.30.3 -google-api-python-client>=2.196.0 -google-auth>=2.53.0 +google-api-core>=2.31.0 +google-api-python-client>=2.198.0 +google-auth>=2.55.1 google-auth-httplib2>=0.4.0 google-auth-oauthlib>=1.4.0 -google-cloud-aiplatform>=1.153.1 -google-cloud-appengine-logging>=1.9.0 -google-cloud-audit-log>=0.5.0 -google-cloud-bigquery>=3.41.0 +google-cloud-aiplatform>=1.159.0 +google-cloud-appengine-logging>=1.10.0 +google-cloud-audit-log>=0.6.0 +google-cloud-bigquery>=3.42.1 google-cloud-core>=2.6.0 -google-cloud-logging>=3.15.0 -google-cloud-mldiagnostics>=1.0.2 -google-cloud-monitoring>=2.30.0 -google-cloud-resource-manager>=1.17.0 -google-cloud-storage>=3.10.1 -google-cloud-storage-control>=1.11.0 +google-cloud-logging>=3.16.0 +google-cloud-mldiagnostics>=1.0.4 +google-cloud-monitoring>=2.31.0 +google-cloud-resource-manager>=1.18.0 +google-cloud-storage>=3.12.0 +google-cloud-storage-control>=1.12.0 google-crc32c>=1.8.0 -google-genai>=2.4.0 +google-genai>=2.10.0 google-pasta>=0.2.0 -google-resumable-media>=2.9.0 +google-resumable-media>=2.10.0 googleapis-common-protos>=1.75.0 -grain>=0.2.16 +grain>=0.2.18 grpc-google-iam-v1>=0.14.4 grpcio>=1.80.0 grpcio-status>=1.80.0 gviz-api>=1.10.0 h11>=0.16.0 h5py>=3.14.0 -hf-xet>=1.5.0 ; platform_machine == 'AMD64' or platform_machine == 'aarch64' or platform_machine == 'amd64' or platform_machine == 'arm64' or platform_machine == 'x86_64' +hf-xet>=1.5.1 ; platform_machine == 'AMD64' or platform_machine == 'aarch64' or platform_machine == 'amd64' or platform_machine == 'arm64' or platform_machine == 'x86_64' httpcore>=1.0.9 -httplib2>=0.31.2 +httplib2>=0.32.0 httpx>=0.28.1 -huggingface-hub>=1.15.0 -humanize>=4.15.0 +huggingface-hub>=1.21.0 +humanize>=4.16.0 hypothesis>=6.142.1 identify>=2.6.19 idna>=3.11 @@ -99,14 +102,16 @@ jax>=0.10.0 jax-cuda12-pjrt>=0.10.0 ; sys_platform == 'linux' jax-cuda12-plugin>=0.10.0 ; sys_platform == 'linux' jaxlib>=0.10.0 -jaxtyping>=0.3.9 +jaxtyping>=0.3.11 jinja2>=3.1.6 +joblib>=1.5.3 jsonlines>=4.0.0 keras>=3.14.0 kiwisolver>=1.5.0 latex2sympy2-extended>=1.11.0 libclang>=18.1.1 libcst>=1.8.6 +lxml>=6.1.1 markdown>=3.10.2 markdown-it-py>=4.0.0 markupsafe>=3.0.3 @@ -116,9 +121,9 @@ mccabe>=0.7.0 mdurl>=0.1.2 ml-collections>=1.1.0 ml-dtypes>=0.5.4 -ml-goodput-measurement>=0.0.16 +ml-goodput-measurement>=0.2.0 mpmath>=1.3.0 -msgpack>=1.1.2 +msgpack>=1.2.1 msgspec>=0.21.1 multidict>=6.7.1 multiprocess>=0.70.19 @@ -127,6 +132,7 @@ namex>=0.1.0 nest-asyncio>=1.6.0 ; sys_platform == 'win32' networkx>=3.6.1 ninja>=1.13.0 +nltk>=3.9.4 nodeenv>=1.10.0 numpy>=2.0.2 numpy-typing-compat>=20251206.2.0 @@ -141,12 +147,13 @@ nvidia-cudnn-cu12>=9.21.0.82 ; sys_platform == 'linux' nvidia-cufft-cu12>=11.4.1.4 ; sys_platform == 'linux' nvidia-cusolver-cu12>=11.7.5.82 ; sys_platform == 'linux' nvidia-cusparse-cu12>=12.5.10.65 ; sys_platform == 'linux' +nvidia-ml-py>=13.610.43 nvidia-nccl-cu12>=2.29.7 ; sys_platform == 'linux' nvidia-nvjitlink-cu12>=12.9.86 ; sys_platform == 'linux' nvidia-nvshmem-cu12>=3.6.5 ; sys_platform == 'linux' oauthlib>=3.3.1 -omegaconf>=2.3.0 -opentelemetry-api>=1.42.0 +omegaconf>=2.3.1 +opentelemetry-api>=1.43.0 opt-einsum>=3.4.0 optax>=0.2.8 optree>=0.19.0 @@ -156,9 +163,9 @@ packaging>=26.1 pandas>=3.0.3 parameterized>=0.9.0 pathspec>=1.1.1 -pathwaysutils>=0.1.8 +pathwaysutils>=0.1.10 pillow>=12.2.0 -platformdirs>=4.9.6 +platformdirs>=4.10.0 pluggy>=1.6.0 portpicker>=1.6.0 pre-commit>=4.6.0 @@ -172,6 +179,7 @@ pyasn1>=0.6.3 pyasn1-modules>=0.4.2 pycnite>=2024.7.31 pycparser>=3.0 ; implementation_name != 'PyPy' and platform_python_implementation != 'PyPy' +pycryptodomex>=3.23.0 pydantic>=2.13.4 pydantic-core>=2.46.4 pydot>=4.0.1 @@ -179,34 +187,37 @@ pyelftools>=0.32 pyglove>=0.4.5 pygments>=2.20.0 pyink>=25.12.0 -pylint>=4.0.5 +pylint>=4.0.6 +pynvml>=13.0.1 +pyopenssl>=26.3.0 pyparsing>=3.3.2 pyproject-hooks>=1.2.0 pytest>=8.4.2 pytest-xdist>=3.8.0 python-dateutil>=2.9.0.post0 -python-discovery>=1.3.1 +python-discovery>=1.4.2 pytokens>=0.4.1 pytype>=2024.10.11 pyyaml>=6.0.3 -qwix>=0.1.6 -regex>=2026.5.9 +qwix>=0.1.8 +regex>=2026.6.28 requests>=2.33.1 requests-oauthlib>=2.0.0 rich>=15.0.0 -safetensors>=0.7.0 +safetensors>=0.8.0 scipy>=1.17.1 scipy-stubs>=1.17.1.4 sentencepiece>=0.2.1 seqio>=0.0.20 setuptools>=82.0.1 shellingham>=1.5.4 +shortuuid>=1.0.13 simple-parsing>=0.1.8 simplejson>=4.1.1 six>=1.17.0 sniffio>=1.3.1 sortedcontainers>=2.4.0 -starlette>=1.0.0 +starlette>=1.3.1 sympy>=1.14.0 tabulate>=0.10.0 tenacity>=9.1.4 @@ -216,7 +227,7 @@ tensorboard-plugin-profile>=2.13.0 tensorboardx>=2.6.5 tensorflow>=2.20.0 tensorflow-datasets>=4.9.10 -tensorflow-metadata>=1.17.3 +tensorflow-metadata>=1.21.0 tensorflow-text>=2.20.1 tensorstore>=0.1.82 termcolor>=3.3.0 @@ -226,11 +237,11 @@ tokenizers>=0.22.2 toml>=0.10.2 tomlkit>=0.15.0 toolz>=1.1.0 -tqdm>=4.66.3 -transformer-engine>=2.15.0 -transformer-engine-cu12>=2.15.0 -transformer-engine-jax>=2.15.0 -transformers>=5.9.0 +tqdm>=4.68.3 +transformer-engine>=2.16.1 +transformer-engine-cu12>=2.16.1 +transformer-engine-jax>=2.16.1 +transformers>=5.12.1 treescope>=0.1.10 typeguard>=2.13.3 typer>=0.25.1 @@ -239,15 +250,15 @@ typing-inspection>=0.4.2 tzdata>=2026.2 ; sys_platform == 'emscripten' or sys_platform == 'win32' uritemplate>=4.2.0 urllib3>=2.6.3 -uvicorn>=0.47.0 +uvicorn>=0.49.0 uvloop>=0.22.1 -virtualenv>=21.3.3 +virtualenv>=21.5.1 wadler-lindig>=0.1.7 websockets>=16.0 werkzeug>=3.1.8 wheel>=0.46.3 -wrapt>=2.1.2 -xxhash>=3.7.0 +wrapt>=2.2.2 +xxhash>=3.8.0 yarl>=1.24.2 zipp>=3.23.1 zstandard>=0.25.0 diff --git a/src/dependencies/requirements/generated_requirements/tpu-requirements.txt b/src/dependencies/requirements/generated_requirements/tpu-requirements.txt index 42eb293874..7f26793d1e 100644 --- a/src/dependencies/requirements/generated_requirements/tpu-requirements.txt +++ b/src/dependencies/requirements/generated_requirements/tpu-requirements.txt @@ -3,7 +3,7 @@ absl-py>=2.4.0 aiofiles>=25.1.0 -aiohappyeyeballs>=2.6.2 +aiohappyeyeballs>=2.7.1 aiohttp>=3.14.1 aiosignal>=1.4.0 annotated-doc>=0.0.4 @@ -15,10 +15,11 @@ array-record>=0.8.3 astroid>=4.0.4 astunparse>=1.6.3 attrs>=26.1.0 -auditwheel>=6.6.0 +auditwheel>=6.7.0 black>=25.12.0 -build>=1.4.3 -certifi>=2026.2.25 +blobfile>=3.2.0 +build>=1.5.0 +certifi>=2026.6.17 cffi>=2.0.0 ; platform_python_implementation != 'PyPy' cfgv>=3.5.0 charset-normalizer>=3.4.7 @@ -29,9 +30,10 @@ cloudpickle>=3.1.2 clu>=0.0.12 colorama>=0.4.6 contourpy>=1.3.3 +coverage>=7.15.0 cryptography>=49.0.0 cycler>=0.12.1 -datasets>=5.0.0 +datasets>=4.8.5 decorator>=5.3.1 dill>=0.4.1 distlib>=0.4.3 @@ -43,28 +45,29 @@ editdistance>=0.8.1 einops>=0.8.2 einshape>=1.0 etils>=1.14.0 +evaluate>=0.2.0 execnet>=2.1.2 -fastapi>=0.138.0 -filelock>=3.28.0 +fastapi>=0.139.0 +filelock>=3.29.4 flatbuffers>=25.12.19 flax>=0.12.7 -fonttools>=4.62.1 +fonttools>=4.63.0 frozenlist>=1.8.0 fsspec>=2026.2.0 gast>=0.7.0 gcsfs>=2026.2.0 google-api-core>=2.31.0 -google-api-python-client>=2.197.0 -google-auth>=2.55.0 +google-api-python-client>=2.198.0 +google-auth>=2.55.1 google-auth-httplib2>=0.4.0 google-auth-oauthlib>=1.4.0 -google-cloud-aiplatform>=1.158.0 +google-cloud-aiplatform>=1.159.0 google-cloud-appengine-logging>=1.10.0 google-cloud-audit-log>=0.6.0 google-cloud-bigquery>=3.42.1 google-cloud-core>=2.6.0 google-cloud-logging>=3.16.0 -google-cloud-mldiagnostics>=1.0.3 +google-cloud-mldiagnostics>=1.0.4 google-cloud-monitoring>=2.31.0 google-cloud-resource-manager>=1.18.0 google-cloud-storage>=3.12.0 @@ -76,40 +79,42 @@ google-resumable-media>=2.10.0 googleapis-common-protos>=1.75.0 grain>=0.2.18 grpc-google-iam-v1>=0.14.4 -grpcio>=1.80.0 -grpcio-status>=1.80.0 +grpcio>=1.81.1 +grpcio-status>=1.81.1 gviz-api>=1.10.0 h11>=0.16.0 h5py>=3.14.0 hf-xet>=1.5.1 ; platform_machine == 'AMD64' or platform_machine == 'aarch64' or platform_machine == 'amd64' or platform_machine == 'arm64' or platform_machine == 'x86_64' httpcore>=1.0.9 -httplib2>=0.31.2 +httplib2>=0.32.0 httpx>=0.28.1 -huggingface-hub>=1.20.1 -humanize>=4.15.0 +huggingface-hub>=1.22.0 +humanize>=4.16.0 hypothesis>=6.142.1 identify>=2.6.19 -idna>=3.11 +idna>=3.18 immutabledict>=4.3.1 importlab>=0.8.1 iniconfig>=2.3.0 isort>=8.0.1 -jax>=0.10.0 -jaxlib>=0.10.0 +jax>=0.10.2 +jaxlib>=0.10.2 jaxtyping>=0.3.11 jinja2>=3.1.6 +joblib>=1.5.3 jsonlines>=4.0.0 -keras>=3.14.0 +keras>=3.14.1 kiwisolver>=1.5.0 latex2sympy2-extended>=1.11.0 libclang>=18.1.1 libcst>=1.8.6 -libtpu>=0.0.40 ; platform_machine == 'x86_64' and sys_platform == 'linux' +libtpu>=0.0.42.1 ; platform_machine == 'x86_64' and sys_platform == 'linux' +lxml>=6.1.1 markdown>=3.10.2 -markdown-it-py>=4.0.0 +markdown-it-py>=4.2.0 markupsafe>=3.0.3 math-verify>=0.9.0 -matplotlib>=3.10.8 +matplotlib>=3.11.0 mccabe>=0.7.0 mdurl>=0.1.2 ml-collections>=1.1.0 @@ -125,30 +130,31 @@ namex>=0.1.0 nest-asyncio>=1.6.0 ; sys_platform == 'win32' networkx>=3.6.1 ninja>=1.13.0 +nltk>=3.9.4 nodeenv>=1.10.0 numpy>=2.0.2 numpy-typing-compat>=20251206.2.0 -nvidia-cuda-cccl>=13.2.75 +nvidia-cuda-cccl>=13.3.3.3.1 ; sys_platform == 'linux' nvidia-ml-py>=13.610.43 oauthlib>=3.3.1 omegaconf>=2.3.1 opentelemetry-api>=1.43.0 opt-einsum>=3.4.0 optax>=0.2.8 -optree>=0.19.0 -optype>=0.17.0 +optree>=0.19.1 +optype>=0.17.1 orbax-checkpoint>=0.12.1 -packaging>=26.1 +packaging>=26.2 pandas>=3.0.3 parameterized>=0.9.0 pathspec>=1.1.1 -pathwaysutils>=0.1.9 +pathwaysutils>=0.1.10 pillow>=12.2.0 platformdirs>=4.10.0 pluggy>=1.6.0 portpicker>=1.6.0 pre-commit>=4.6.0 -prometheus-client>=0.20.0 +prometheus-client>=0.25.0 promise>=2.3 propcache>=0.5.2 proto-plus>=1.28.0 @@ -159,10 +165,11 @@ pyasn1>=0.6.3 pyasn1-modules>=0.4.2 pycnite>=2024.7.31 pycparser>=3.0 ; implementation_name != 'PyPy' and platform_python_implementation != 'PyPy' +pycryptodomex>=3.23.0 pydantic>=2.13.4 pydantic-core>=2.46.4 pydot>=4.0.1 -pyelftools>=0.32 +pyelftools>=0.33 pyglove>=0.4.5 pygments>=2.20.0 pyink>=25.12.0 @@ -171,25 +178,27 @@ pynvml>=13.0.1 pyopenssl>=26.3.0 pyparsing>=3.3.2 pyproject-hooks>=1.2.0 -pytest>=8.4.2 +pytest>=9.1.0 pytest-xdist>=3.8.0 python-dateutil>=2.9.0.post0 -python-discovery>=1.4.2 +python-discovery>=1.4.3 pytokens>=0.4.1 pytype>=2024.10.11 pyyaml>=6.0.3 qwix>=0.1.8 -regex>=2026.5.9 -requests>=2.33.1 +regex>=2026.6.28 +requests>=2.34.2 requests-oauthlib>=2.0.0 +responses>=0.18.0 rich>=15.0.0 safetensors>=0.8.0 scipy>=1.17.1 -scipy-stubs>=1.17.1.4 +scipy-stubs>=1.17.1.5 sentencepiece>=0.2.1 seqio>=0.0.20 setuptools>=82.0.1 shellingham>=1.5.4 +shortuuid>=1.0.13 simple-parsing>=0.1.8 simplejson>=4.1.1 six>=1.17.0 @@ -216,24 +225,24 @@ toml>=0.10.2 tomlkit>=0.15.0 toolz>=1.1.0 tqdm>=4.68.3 -transformers>=5.12.1 +transformers>=5.13.0 treescope>=0.1.10 typeguard>=2.13.3 -typer>=0.25.1 +typer>=0.26.8 typing-extensions>=4.15.0 typing-inspection>=0.4.2 tzdata>=2026.2 ; sys_platform == 'emscripten' or sys_platform == 'win32' uritemplate>=4.2.0 -urllib3>=2.6.3 -uvicorn>=0.49.0 +urllib3>=2.7.0 +uvicorn>=0.50.2 uvloop>=0.22.1 -virtualenv>=21.5.1 +virtualenv>=21.5.2 wadler-lindig>=0.1.7 websockets>=16.0 werkzeug>=3.1.8 -wheel>=0.46.3 +wheel>=0.47.0 wrapt>=2.2.2 -xxhash>=3.7.1 +xxhash>=3.8.1 yarl>=1.24.2 -zipp>=3.23.1 +zipp>=4.1.0 zstandard>=0.25.0 diff --git a/src/maxtext/checkpoint_conversion/inspect_checkpoint.py b/src/maxtext/checkpoint_conversion/inspect_checkpoint.py index c63f2e1161..7e9784e516 100644 --- a/src/maxtext/checkpoint_conversion/inspect_checkpoint.py +++ b/src/maxtext/checkpoint_conversion/inspect_checkpoint.py @@ -79,7 +79,7 @@ def print_structure(data_dict, output_file=""): """Utility to format and print sorted keys and shapes from a flattened dictionary.""" if output_file: # Save command - save_lines = [f"# {" ".join(sys.orig_argv)}", ""] + save_lines = [f"# {' '.join(sys.orig_argv)}", ""] for key in sorted(data_dict.keys(), key=natural_sort_key): line = f"key: {key} | {data_dict[key]}" diff --git a/src/maxtext/checkpoint_conversion/standalone_scripts/convert_gpt3_ckpt_from_paxml.py b/src/maxtext/checkpoint_conversion/standalone_scripts/convert_gpt3_ckpt_from_paxml.py index b8a513c1df..12bc48c676 100644 --- a/src/maxtext/checkpoint_conversion/standalone_scripts/convert_gpt3_ckpt_from_paxml.py +++ b/src/maxtext/checkpoint_conversion/standalone_scripts/convert_gpt3_ckpt_from_paxml.py @@ -117,7 +117,7 @@ def init_state_fn(): cfg.checkpoint_period, ) - state, _, _, _ = maxtext_utils.setup_training_state(None, cfg, mesh, checkpoint_manager, init_state_fn) + state, _, _, _, _ = maxtext_utils.setup_training_state(None, cfg, mesh, checkpoint_manager, init_state_fn) max_logging.log("start") max_utils.print_mem_stats("After params initialized") diff --git a/src/maxtext/checkpoint_conversion/to_huggingface.py b/src/maxtext/checkpoint_conversion/to_huggingface.py index c52bd8192d..f0e94d4829 100644 --- a/src/maxtext/checkpoint_conversion/to_huggingface.py +++ b/src/maxtext/checkpoint_conversion/to_huggingface.py @@ -97,6 +97,8 @@ from maxtext.utils import max_logging from maxtext.utils import max_utils from maxtext.utils.globals import HF_IDS +from maxtext.utils.lora_utils import sync_lora_metadata +from maxtext.utils.model_creation_utils import verify_and_sync_scan_layers flags.DEFINE_bool( @@ -442,6 +444,14 @@ def main(argv: Sequence[str]) -> None: assert ( config.load_full_state_path == "" ), "This script expects parameters, not a full state. Use generate_param_only_checkpoint first if needed." + + if not config.load_parameters_path and config.lora.lora_restore_path: + # Standalone LoRA conversion uses lora_restore_path for metadata + temp_config = config.model_copy(update={"load_parameters_path": config.lora.lora_restore_path}) + temp_config = verify_and_sync_scan_layers(temp_config) + config = config.model_copy(update={"scan_layers": temp_config.scan_layers}) + else: + config = verify_and_sync_scan_layers(config) max_utils.print_system_information() overall_start = time.time() @@ -451,6 +461,9 @@ def main(argv: Sequence[str]) -> None: if not load_parameters_path and not lora_restore_path: raise ValueError("Either load_parameters_path or lora_restore_path must be specified.") + if lora_restore_path: + sync_lora_metadata(config) + # Load Maxtext checkpoint using Orbax (now smart enough to load both if present) max_logging.log("\nLoading Orbax checkpoint(s)...") start = time.time() diff --git a/src/maxtext/checkpoint_conversion/to_maxtext.py b/src/maxtext/checkpoint_conversion/to_maxtext.py index c2ff040979..ca50d67680 100644 --- a/src/maxtext/checkpoint_conversion/to_maxtext.py +++ b/src/maxtext/checkpoint_conversion/to_maxtext.py @@ -1070,6 +1070,7 @@ def _eager_getter(key): simulated_cpu_devices_count, config.checkpoint_storage_use_ocdbt, config.checkpoint_storage_use_zarr3, + config=config, ) print_ram_usage("Program Ends") diff --git a/src/maxtext/checkpoint_conversion/utils/hf_model_configs.py b/src/maxtext/checkpoint_conversion/utils/hf_model_configs.py index 8bb262958c..c72afdebc7 100644 --- a/src/maxtext/checkpoint_conversion/utils/hf_model_configs.py +++ b/src/maxtext/checkpoint_conversion/utils/hf_model_configs.py @@ -19,7 +19,7 @@ import transformers -if transformers.__version__ >= "5.0.0": +if transformers.__version__ >= "5.0.0": # pyrefly: ignore[missing-attribute] from transformers.configuration_utils import PreTrainedConfig as PTConfig # pytype: disable=import-error else: from transformers.configuration_utils import PretrainedConfig as PTConfig @@ -120,7 +120,7 @@ gemma4_31b_dict = gemma4_26b_dict.copy() gemma4_31b_dict["text_config"] = gemma4_26b_dict["text_config"].copy() -gemma4_31b_dict["text_config"].update( +gemma4_31b_dict["text_config"].update( # pyrefly: ignore[no-matching-overload] { "enable_moe_block": False, "hidden_size": 5376, @@ -265,10 +265,10 @@ try: # Will execute successfully if Transformers is updated with Gemma 4 support - gemma4_26b_config = transformers.Gemma4Config(**gemma4_26b_dict) - gemma4_31b_config = transformers.Gemma4Config(**gemma4_31b_dict) - gemma4_e2b_config = transformers.Gemma4Config(**gemma4_e2b_dict) - gemma4_e4b_config = transformers.Gemma4Config(**gemma4_e4b_dict) + gemma4_26b_config = transformers.Gemma4Config(**gemma4_26b_dict) # pyrefly: ignore[missing-attribute] + gemma4_31b_config = transformers.Gemma4Config(**gemma4_31b_dict) # pyrefly: ignore[missing-attribute] + gemma4_e2b_config = transformers.Gemma4Config(**gemma4_e2b_dict) # pyrefly: ignore[missing-attribute] + gemma4_e4b_config = transformers.Gemma4Config(**gemma4_e4b_dict) # pyrefly: ignore[missing-attribute] except AttributeError: # Graceful fallback to raw dict-based PTConfig if Gemma 4 natively is missing gemma4_26b_config = PTConfig(**gemma4_26b_dict) # pytype: disable=wrong-arg-types @@ -1011,7 +1011,7 @@ # TODO(shuningjin): replace with DeepseekV32Config when available in transformers library -class DeepseekV32Config(PTConfig): +class DeepseekV32Config(PTConfig): # pyrefly: ignore[invalid-inheritance] model_type = "deepseek_v32" def __init__(self, **kwargs): @@ -1089,7 +1089,7 @@ def __init__(self, **kwargs): "use_cache": True, "vocab_size": 201088, } -gpt_oss_20b_config = transformers.GptOssConfig(**gpt_oss_20b_dict) +gpt_oss_20b_config = transformers.GptOssConfig(**gpt_oss_20b_dict) # pyrefly: ignore[bad-argument-type] # from https://huggingface.co/openai/gpt-oss-120b/blob/main/config.json # remove mxfp4 quantization_config, since we are using bf16 @@ -1171,7 +1171,7 @@ def __init__(self, **kwargs): "use_cache": True, "vocab_size": 201088, } -gpt_oss_120b_config = transformers.GptOssConfig(**gpt_oss_120b_dict) +gpt_oss_120b_config = transformers.GptOssConfig(**gpt_oss_120b_dict) # pyrefly: ignore[bad-argument-type] qwen3_omni_30b_a3b_config = transformers.Qwen3OmniMoeConfig( @@ -1481,8 +1481,8 @@ def __init__(self, **kwargs): try: # Will execute successfully if Transformers is updated with Qwen3.5 support - qwen3_5_35b_a3b_config = transformers.Qwen3_5MoeConfig(**qwen3_5_35b_a3b_dict) - qwen3_5_397b_a17b_config = transformers.Qwen3_5MoeConfig(**qwen3_5_397b_a17b_dict) + qwen3_5_35b_a3b_config = transformers.Qwen3_5MoeConfig(**qwen3_5_35b_a3b_dict) # pyrefly: ignore[missing-attribute] + qwen3_5_397b_a17b_config = transformers.Qwen3_5MoeConfig(**qwen3_5_397b_a17b_dict) # pyrefly: ignore[missing-attribute] except AttributeError: qwen3_5_35b_a3b_config = PTConfig(**qwen3_5_35b_a3b_dict) # pytype: disable=wrong-arg-types qwen3_5_397b_a17b_config = PTConfig(**qwen3_5_397b_a17b_dict) # pytype: disable=wrong-arg-types diff --git a/src/maxtext/checkpoint_conversion/utils/param_mapping.py b/src/maxtext/checkpoint_conversion/utils/param_mapping.py index 80707a2abd..be0c24baea 100644 --- a/src/maxtext/checkpoint_conversion/utils/param_mapping.py +++ b/src/maxtext/checkpoint_conversion/utils/param_mapping.py @@ -153,7 +153,7 @@ def GEMMA3_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=False hf_indices = list(range(block_idx, num_scanned, attention_pattern_length)) for mx, hf in text_params: key = f"params-decoder-layers-layers_{block_idx}-{mx}" - mapping[key] = [f"model.language_model.layers.{i}.{hf}" for i in hf_indices] + mapping[key] = [f"model.language_model.layers.{i}.{hf}" for i in hf_indices] # pyrefly: ignore[bad-assignment] # Remainder layers (unscanned): params-decoder-layers_remainder-layers_{rem_idx}-{param} if num_remaining > 0: @@ -623,7 +623,7 @@ def QWEN_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=False): if scan_layers: # This block handles scanned layers for both dense and MoE models. - mapping.update( + mapping.update( # pyrefly: ignore[no-matching-overload] { "params-decoder-layers-pre_self_attention_layer_norm-scale": [ f"model.layers.{i}.input_layernorm.weight" for i in range(n_layers) @@ -663,7 +663,7 @@ def QWEN_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=False): if num_experts > 1: # For scanned MoE, we create a nested list: [[e0_l0, e0_l1..], [e1_l0, e1_l1..]..] # This follows the (experts, layers, ...) tensor layout. - mapping.update( + mapping.update( # pyrefly: ignore[no-matching-overload] { "params-decoder-layers-moe_block-gate-kernel": [ f"model.layers.{i}.mlp.gate.weight" for i in range(n_layers) @@ -683,7 +683,7 @@ def QWEN_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=False): } ) else: # Dense MLP - mapping.update( + mapping.update( # pyrefly: ignore[no-matching-overload] { "params-decoder-layers-mlp-wi_0-kernel": [ f"model.layers.{i}.mlp.gate_proj.weight" for i in range(n_layers) @@ -714,7 +714,7 @@ def QWEN_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=False): ) if num_experts > 1: # For each unscanned MoE layer, map the MaxText parameter to a 1D list of all expert weights for that layer. - mapping.update( + mapping.update( # pyrefly: ignore[no-matching-overload] { f"params-decoder-layers_{i}-moe_block-gate-kernel": f"model.layers.{i}.mlp.gate.weight", f"params-decoder-layers_{i}-moe_block-wi_0": [ @@ -871,10 +871,10 @@ def QWEN3_5_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=Fals prefix = f"params-decoder-layers-layer_{block_idx}" # Layer norms - mapping[f"{prefix}-input_layernorm-scale"] = [ + mapping[f"{prefix}-input_layernorm-scale"] = [ # pyrefly: ignore[bad-assignment] f"model.language_model.layers.{i}.input_layernorm.weight" for i in hf_indices ] - mapping[f"{prefix}-post_attention_layernorm-scale"] = [ + mapping[f"{prefix}-post_attention_layernorm-scale"] = [ # pyrefly: ignore[bad-assignment] f"model.language_model.layers.{i}.post_attention_layernorm.weight" for i in hf_indices ] @@ -882,7 +882,7 @@ def QWEN3_5_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=Fals is_full_attention_layer = (block_idx + 1) % layer_cycle_interval == 0 if is_full_attention_layer: - mapping.update( + mapping.update( # pyrefly: ignore[no-matching-overload] { f"{prefix}-attention-attention-query-kernel": [ f"model.language_model.layers.{i}.self_attn.q_proj.weight" for i in hf_indices @@ -906,7 +906,7 @@ def QWEN3_5_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=Fals ) else: # Linear/Hybrid Attention Block - mapping.update( + mapping.update( # pyrefly: ignore[no-matching-overload] { # Provide a tuple of HF keys so MaxText concatenates them into qkvz f"{prefix}-attention-in_proj_qkvz-kernel": [ @@ -941,7 +941,7 @@ def QWEN3_5_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=Fals ) # 3. Handle MLP: Gates and Shared Experts - mapping.update( + mapping.update( # pyrefly: ignore[no-matching-overload] { f"{prefix}-mlp-routed_experts-gate-kernel": [ f"model.language_model.layers.{i}.mlp.gate.weight" for i in hf_indices @@ -962,7 +962,7 @@ def QWEN3_5_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=Fals ) # 4. Handle MoE Routed Experts - mapping.update( + mapping.update( # pyrefly: ignore[no-matching-overload] { f"{prefix}-mlp-routed_experts-wo": [ f"model.language_model.layers.{i}.mlp.experts.down_proj" for i in hf_indices @@ -999,7 +999,7 @@ def QWEN3_5_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=Fals ) else: # Linear/Hybrid Attention Block (Unscanned) - mapping.update( + mapping.update( # pyrefly: ignore[no-matching-overload] { # Provide a tuple of HF keys so MaxText concatenates them into qkvz f"{prefix}-attention-in_proj_qkvz-kernel": ( @@ -1033,7 +1033,7 @@ def QWEN3_5_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=Fals ) # MoE Routed Experts - mapping.update( + mapping.update( # pyrefly: ignore[no-matching-overload] { f"{prefix}-mlp-routed_experts-wo": f"model.language_model.layers.{i}.mlp.experts.down_proj", ( @@ -1255,7 +1255,7 @@ def concat_ba_and_transpose(input_tensor, target_shape=None): if is_full_attention_layer: for key in ["query", "key", "value", "out"]: - hooks[f"{prefix}-attention-attention-{key}-kernel"] = reshape_kernel + hooks[f"{prefix}-attention-attention-{key}-kernel"] = reshape_kernel # pyrefly: ignore[bad-assignment] else: hooks[f"{prefix}-attention-in_proj_qkvz-kernel"] = concat_qkvz_and_transpose hooks[f"{prefix}-attention-in_proj_ba-kernel"] = concat_ba_and_transpose @@ -1269,7 +1269,7 @@ def concat_ba_and_transpose(input_tensor, target_shape=None): hooks[f"{mlp_prefix}-shared_expert-wo-kernel"] = transpose hooks[f"{mlp_prefix}-shared_expert_gate-kernel"] = transpose - hooks[(f"{mlp_prefix}-routed_experts-wi_0", f"{mlp_prefix}-routed_experts-wi_1")] = process_wi_0_wi_1 + hooks[(f"{mlp_prefix}-routed_experts-wi_0", f"{mlp_prefix}-routed_experts-wi_1")] = process_wi_0_wi_1 # pyrefly: ignore[unsupported-operation] hooks[f"{mlp_prefix}-routed_experts-wo"] = transpose_expert # Vision hooks for Qwen3.5 @@ -1340,23 +1340,23 @@ def reshape_vision_attn_out(input_tensor, target_shape): return input_tensor.T.reshape(target_shape) # Apply vision hooks - hooks["params-vision_encoder-Qwen3_5MoeVisionEncoder_0-patch_embed-proj-kernel"] = reshape_conv3d_patch_embed + hooks["params-vision_encoder-Qwen3_5MoeVisionEncoder_0-patch_embed-proj-kernel"] = reshape_conv3d_patch_embed # pyrefly: ignore[bad-assignment] for i in range(n_vision_layers): prefix = f"params-vision_encoder-Qwen3_5MoeVisionEncoder_0-blocks_{i}" - hooks[f"{prefix}-attn-attn-query-kernel"] = split_qkv_query - hooks[f"{prefix}-attn-attn-query-bias"] = split_qkv_bias_query - hooks[f"{prefix}-attn-attn-key-kernel"] = split_qkv_key - hooks[f"{prefix}-attn-attn-key-bias"] = split_qkv_bias_key - hooks[f"{prefix}-attn-attn-value-kernel"] = split_qkv_value - hooks[f"{prefix}-attn-attn-value-bias"] = split_qkv_bias_value - hooks[f"{prefix}-attn-attn-out-kernel"] = reshape_vision_attn_out - hooks[f"{prefix}-mlp-kernel"] = reshape_kernel_vision - hooks[f"{prefix}-mlp_out-kernel"] = reshape_kernel_vision + hooks[f"{prefix}-attn-attn-query-kernel"] = split_qkv_query # pyrefly: ignore[bad-assignment] + hooks[f"{prefix}-attn-attn-query-bias"] = split_qkv_bias_query # pyrefly: ignore[bad-assignment] + hooks[f"{prefix}-attn-attn-key-kernel"] = split_qkv_key # pyrefly: ignore[bad-assignment] + hooks[f"{prefix}-attn-attn-key-bias"] = split_qkv_bias_key # pyrefly: ignore[bad-assignment] + hooks[f"{prefix}-attn-attn-value-kernel"] = split_qkv_value # pyrefly: ignore[bad-assignment] + hooks[f"{prefix}-attn-attn-value-bias"] = split_qkv_bias_value # pyrefly: ignore[bad-assignment] + hooks[f"{prefix}-attn-attn-out-kernel"] = reshape_vision_attn_out # pyrefly: ignore[bad-assignment] + hooks[f"{prefix}-mlp-kernel"] = reshape_kernel_vision # pyrefly: ignore[bad-assignment] + hooks[f"{prefix}-mlp_out-kernel"] = reshape_kernel_vision # pyrefly: ignore[bad-assignment] # Vision projector - hooks["params-vision_encoder-Qwen3_5MoeVisionProjector_0-merger-mlp_0-kernel"] = reshape_kernel_vision - hooks["params-vision_encoder-Qwen3_5MoeVisionProjector_0-merger-mlp_2-kernel"] = reshape_kernel_vision + hooks["params-vision_encoder-Qwen3_5MoeVisionProjector_0-merger-mlp_0-kernel"] = reshape_kernel_vision # pyrefly: ignore[bad-assignment] + hooks["params-vision_encoder-Qwen3_5MoeVisionProjector_0-merger-mlp_2-kernel"] = reshape_kernel_vision # pyrefly: ignore[bad-assignment] return hooks @@ -1384,8 +1384,8 @@ def QWEN3_NEXT_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=F prefix = f"params-decoder-layers-layer_{block_idx}" # Layer norms - mapping[f"{prefix}-input_layernorm-scale"] = [f"model.layers.{i}.input_layernorm.weight" for i in hf_indices] - mapping[f"{prefix}-post_attention_layernorm-scale"] = [ + mapping[f"{prefix}-input_layernorm-scale"] = [f"model.layers.{i}.input_layernorm.weight" for i in hf_indices] # pyrefly: ignore[bad-assignment] + mapping[f"{prefix}-post_attention_layernorm-scale"] = [ # pyrefly: ignore[bad-assignment] f"model.layers.{i}.post_attention_layernorm.weight" for i in hf_indices ] @@ -1393,7 +1393,7 @@ def QWEN3_NEXT_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=F is_full_attention_layer = (block_idx + 1) % layer_cycle_interval == 0 if is_full_attention_layer: - mapping.update( + mapping.update( # pyrefly: ignore[no-matching-overload] { f"{prefix}-attention-attention-query-kernel": [ f"model.layers.{i}.self_attn.q_proj.weight" for i in hf_indices @@ -1417,7 +1417,7 @@ def QWEN3_NEXT_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=F ) else: # Linear/Hybrid Attention Block - mapping.update( + mapping.update( # pyrefly: ignore[no-matching-overload] { f"{prefix}-attention-in_proj_qkvz-kernel": [ f"model.layers.{i}.linear_attn.in_proj_qkvz.weight" for i in hf_indices @@ -1438,7 +1438,7 @@ def QWEN3_NEXT_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=F ) # 3. Handle MLP: Gates and Shared Experts - mapping.update( + mapping.update( # pyrefly: ignore[no-matching-overload] { f"{prefix}-mlp-routed_experts-gate-kernel": [f"model.layers.{i}.mlp.gate.weight" for i in hf_indices], f"{prefix}-mlp-shared_expert-wi_0-kernel": [ @@ -1457,7 +1457,7 @@ def QWEN3_NEXT_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=F ) # 4. Handle MoE Routed Experts - mapping.update( + mapping.update( # pyrefly: ignore[no-matching-overload] { f"{prefix}-mlp-routed_experts-wi_0": [ [f"model.layers.{i}.mlp.experts.{e}.gate_proj.weight" for i in hf_indices] for e in range(num_experts) @@ -1521,7 +1521,7 @@ def QWEN3_NEXT_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=F ) # MoE Routed Experts (List of expert weights for this specific layer) - mapping.update( + mapping.update( # pyrefly: ignore[no-matching-overload] { f"{prefix}-mlp-routed_experts-wi_0": [ f"model.layers.{i}.mlp.experts.{e}.gate_proj.weight" for e in range(num_experts) @@ -1576,7 +1576,7 @@ def permute_conv(input_tensor, target_shape=None): if is_full_attention_layer: for key in ["query", "key", "value", "out"]: - hooks[f"{prefix}-attention-attention-{key}-kernel"] = reshape_kernel + hooks[f"{prefix}-attention-attention-{key}-kernel"] = reshape_kernel # pyrefly: ignore[bad-assignment] else: hooks[f"{prefix}-attention-in_proj_qkvz-kernel"] = transpose hooks[f"{prefix}-attention-in_proj_ba-kernel"] = transpose @@ -1663,17 +1663,17 @@ def DEEPSEEK_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=Fal # scan if scan_layers: for maxtext_key, hf_key in dense_layer_keys.items(): - mapping[f"params-decoder-dense_layers-{maxtext_key}"] = [ + mapping[f"params-decoder-dense_layers-{maxtext_key}"] = [ # pyrefly: ignore[bad-assignment] f"model.layers.{i}.{hf_key}" for i in range(first_num_dense_layers) ] for maxtext_key, hf_key in moe_layer_keys.items(): - mapping[f"params-decoder-moe_layers-{maxtext_key}"] = [ + mapping[f"params-decoder-moe_layers-{maxtext_key}"] = [ # pyrefly: ignore[bad-assignment] f"model.layers.{i}.{hf_key}" for i in range(first_num_dense_layers, num_main_layers) ] for maxtext_key, hf_key in moe_expert_keys.items(): - mapping[f"params-decoder-moe_layers-{maxtext_key}"] = [ + mapping[f"params-decoder-moe_layers-{maxtext_key}"] = [ # pyrefly: ignore[bad-assignment] [f"model.layers.{i}.mlp.experts.{e}.{hf_key}" for i in range(first_num_dense_layers, num_main_layers)] for e in range(num_experts) ] @@ -1690,7 +1690,7 @@ def DEEPSEEK_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=Fal mapping[f"params-decoder-moe_layers_{moe_layer_idx}-{maxtext_key}"] = f"model.layers.{i}.{hf_key}" for maxtext_key, hf_key in moe_expert_keys.items(): - mapping[f"params-decoder-moe_layers_{moe_layer_idx}-{maxtext_key}"] = [ + mapping[f"params-decoder-moe_layers_{moe_layer_idx}-{maxtext_key}"] = [ # pyrefly: ignore[bad-assignment] f"model.layers.{i}.mlp.experts.{e}.{hf_key}" for e in range(num_experts) ] return mapping @@ -1833,7 +1833,7 @@ def GPT_OSS_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=Fals f"model.layers.{i}.mlp.experts.gate_up_proj_bias" for i in hf_indices ], } - mapping.update(block_mapping) + mapping.update(block_mapping) # pyrefly: ignore[no-matching-overload] else: # Unscan @@ -1868,7 +1868,7 @@ def GPT_OSS_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=Fals f"{prefix}-GptOssMlp-wi_1_bias", ): f"model.layers.{i}.mlp.experts.gate_up_proj_bias", } - mapping.update(layer_mapping) + mapping.update(layer_mapping) # pyrefly: ignore[no-matching-overload] return mapping @@ -1916,7 +1916,7 @@ def interleave(input_tensor, target_shape=None): """ if saving_to_hf: wi_0, wi_1 = input_tensor - wi_0_1 = np.empty(target_shape, dtype=wi_0.dtype) + wi_0_1 = np.empty(target_shape, dtype=wi_0.dtype) # pyrefly: ignore[no-matching-overload] wi_0_1[..., ::2] = wi_0 wi_0_1[..., 1::2] = wi_1 return wi_0_1 @@ -1936,14 +1936,14 @@ def interleave(input_tensor, target_shape=None): prefix = f"params-decoder-layers-layers_{idx}" if scan_layers else f"params-decoder-layers_{idx}" # Attention Kernels & Biases for key in ["query", "key", "value"]: - hooks[f"{prefix}-GptOssAttention-{key}-kernel"] = reshape_kernel + hooks[f"{prefix}-GptOssAttention-{key}-kernel"] = reshape_kernel # pyrefly: ignore[bad-assignment] hooks[f"{prefix}-GptOssAttention-{key}-bias"] = reshape_bias - hooks[f"{prefix}-GptOssAttention-out-kernel"] = reshape_kernel + hooks[f"{prefix}-GptOssAttention-out-kernel"] = reshape_kernel # pyrefly: ignore[bad-assignment] # MLP Kernels & Biases hooks[f"{prefix}-GptOssMlp-gate-kernel"] = transpose # `composite_mt_key`: A hook for combining multiple MaxText params. - hooks[(f"{prefix}-GptOssMlp-wi_0", f"{prefix}-GptOssMlp-wi_1")] = interleave - hooks[(f"{prefix}-GptOssMlp-wi_0_bias", f"{prefix}-GptOssMlp-wi_1_bias")] = interleave + hooks[(f"{prefix}-GptOssMlp-wi_0", f"{prefix}-GptOssMlp-wi_1")] = interleave # pyrefly: ignore[unsupported-operation] + hooks[(f"{prefix}-GptOssMlp-wi_0_bias", f"{prefix}-GptOssMlp-wi_1_bias")] = interleave # pyrefly: ignore[unsupported-operation] return hooks @@ -2407,31 +2407,31 @@ def LLAMA31_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=Fals } if scan_layers: - mapping["params-decoder-layers-self_attention-query-kernel"] = [ + mapping["params-decoder-layers-self_attention-query-kernel"] = [ # pyrefly: ignore[bad-assignment] f"model.layers.{layer_idx}.self_attn.q_proj.weight" for layer_idx in range(n_layers) ] - mapping["params-decoder-layers-self_attention-key-kernel"] = [ + mapping["params-decoder-layers-self_attention-key-kernel"] = [ # pyrefly: ignore[bad-assignment] f"model.layers.{layer_idx}.self_attn.k_proj.weight" for layer_idx in range(n_layers) ] - mapping["params-decoder-layers-self_attention-value-kernel"] = [ + mapping["params-decoder-layers-self_attention-value-kernel"] = [ # pyrefly: ignore[bad-assignment] f"model.layers.{layer_idx}.self_attn.v_proj.weight" for layer_idx in range(n_layers) ] - mapping["params-decoder-layers-self_attention-out-kernel"] = [ + mapping["params-decoder-layers-self_attention-out-kernel"] = [ # pyrefly: ignore[bad-assignment] f"model.layers.{layer_idx}.self_attn.o_proj.weight" for layer_idx in range(n_layers) ] - mapping["params-decoder-layers-mlp-wi_0-kernel"] = [ + mapping["params-decoder-layers-mlp-wi_0-kernel"] = [ # pyrefly: ignore[bad-assignment] f"model.layers.{layer_idx}.mlp.gate_proj.weight" for layer_idx in range(n_layers) ] - mapping["params-decoder-layers-mlp-wi_1-kernel"] = [ + mapping["params-decoder-layers-mlp-wi_1-kernel"] = [ # pyrefly: ignore[bad-assignment] f"model.layers.{layer_idx}.mlp.up_proj.weight" for layer_idx in range(n_layers) ] - mapping["params-decoder-layers-mlp-wo-kernel"] = [ + mapping["params-decoder-layers-mlp-wo-kernel"] = [ # pyrefly: ignore[bad-assignment] f"model.layers.{layer_idx}.mlp.down_proj.weight" for layer_idx in range(n_layers) ] - mapping["params-decoder-layers-pre_self_attention_layer_norm-scale"] = [ + mapping["params-decoder-layers-pre_self_attention_layer_norm-scale"] = [ # pyrefly: ignore[bad-assignment] f"model.layers.{layer_idx}.input_layernorm.weight" for layer_idx in range(n_layers) ] - mapping["params-decoder-layers-post_self_attention_layer_norm-scale"] = [ + mapping["params-decoder-layers-post_self_attention_layer_norm-scale"] = [ # pyrefly: ignore[bad-assignment] f"model.layers.{layer_idx}.post_attention_layernorm.weight" for layer_idx in range(n_layers) ] else: @@ -2816,7 +2816,7 @@ def GEMMA4_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=False for layer_in_block in range(attention_pattern_length): hf_indices = list(range(layer_in_block, num_scanned, attention_pattern_length)) prefix = f"params-decoder-scanned_blocks-layers_{layer_in_block}" - mapping.update( + mapping.update( # pyrefly: ignore[no-matching-overload] { f"{prefix}-self_attention-query-kernel": [ f"{text_base}.layers.{i}.self_attn.q_proj.weight" for i in hf_indices @@ -2841,14 +2841,14 @@ def GEMMA4_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=False } ) if maxtext_config.v_norm_with_scale: - mapping.update( + mapping.update( # pyrefly: ignore[no-matching-overload] { f"{prefix}-self_attention-value_norm-scale": [ f"{text_base}.layers.{i}.self_attn.v_norm.weight" for i in hf_indices ] } ) - mapping.update( + mapping.update( # pyrefly: ignore[no-matching-overload] { f"{prefix}-pre_self_attention_norm-scale": [ f"{text_base}.layers.{i}.input_layernorm.weight" for i in hf_indices @@ -2946,7 +2946,7 @@ def GEMMA4_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=False ) if maxtext_config.v_norm_with_scale: mapping.update({f"{prefix}-self_attention-value_norm-scale": (f"{hf_prefix}.self_attn.v_norm.weight")}) - mapping.update( + mapping.update( # pyrefly: ignore[no-matching-overload] { f"{prefix}-pre_self_attention_norm-scale": (f"{hf_prefix}.input_layernorm.weight"), f"{prefix}-post_self_attention_norm-scale": (f"{hf_prefix}.post_attention_layernorm.weight"), @@ -2998,7 +2998,7 @@ def GEMMA4_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=False is_global = i % 6 == 5 num_experts = tcfg.get("num_experts") num_experts = num_experts if num_experts is not None else 1 - mapping.update( + mapping.update( # pyrefly: ignore[no-matching-overload] { f"{prefix}-self_attention-query-kernel": (f"{hf_prefix}.self_attn.q_proj.weight"), f"{prefix}-self_attention-key-kernel": (f"{hf_prefix}.self_attn.k_proj.weight"), @@ -3012,7 +3012,7 @@ def GEMMA4_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=False ) if maxtext_config.v_norm_with_scale: mapping.update({f"{prefix}-self_attention-value_norm-scale": (f"{hf_prefix}.self_attn.v_norm.weight")}) - mapping.update( + mapping.update( # pyrefly: ignore[no-matching-overload] { f"{prefix}-pre_self_attention_norm-scale": (f"{hf_prefix}.input_layernorm.weight"), f"{prefix}-post_self_attention_norm-scale": (f"{hf_prefix}.post_attention_layernorm.weight"), @@ -3111,7 +3111,7 @@ def GEMMA4_SMALL_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers } ) if not is_shared: - mapping.update( + mapping.update( # pyrefly: ignore[no-matching-overload] { f"{prefix}-self_attention-key-kernel": f"{hf_prefix}.self_attn.k_proj.weight", f"{prefix}-self_attention-value-kernel": f"{hf_prefix}.self_attn.v_proj.weight", @@ -3527,7 +3527,7 @@ def OLMO3_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=False) hf_indices = range(cycle_idx, n_layers, layer_cycle_interval) prefix = f"params-decoder-layers-layers_{cycle_idx}" - mapping.update( + mapping.update( # pyrefly: ignore[no-matching-overload] { # Attention Projections f"{prefix}-attention-query-kernel": [f"model.layers.{i}.self_attn.q_proj.weight" for i in hf_indices], @@ -3791,9 +3791,9 @@ def process_qkv_vision(input_tensor, target_shape=None): q_hf = input_tensor[:hidden_size, :] k_hf = input_tensor[hidden_size : 2 * hidden_size, :] v_hf = input_tensor[2 * hidden_size :, :] - q_mt = q_hf.T.reshape(target_shape[0]) - k_mt = k_hf.T.reshape(target_shape[1]) - v_mt = v_hf.T.reshape(target_shape[2]) + q_mt = q_hf.T.reshape(target_shape[0]) # pyrefly: ignore[unsupported-operation] + k_mt = k_hf.T.reshape(target_shape[1]) # pyrefly: ignore[unsupported-operation] + v_mt = v_hf.T.reshape(target_shape[2]) # pyrefly: ignore[unsupported-operation] return np.stack([q_mt, k_mt, v_mt], axis=-1) def process_qkv_bias_vision(input_tensor, target_shape=None): @@ -3808,9 +3808,9 @@ def process_qkv_bias_vision(input_tensor, target_shape=None): qb_hf = input_tensor[:hidden_size] kb_hf = input_tensor[hidden_size : 2 * hidden_size] vb_hf = input_tensor[2 * hidden_size :] - qb_mt = qb_hf.reshape(target_shape[0]) - kb_mt = kb_hf.reshape(target_shape[1]) - vb_mt = vb_hf.reshape(target_shape[2]) + qb_mt = qb_hf.reshape(target_shape[0]) # pyrefly: ignore[unsupported-operation] + kb_mt = kb_hf.reshape(target_shape[1]) # pyrefly: ignore[unsupported-operation] + vb_mt = vb_hf.reshape(target_shape[2]) # pyrefly: ignore[unsupported-operation] return np.stack([qb_mt, kb_mt, vb_mt], axis=-1) def reshape_vision_attn_out(input_tensor, target_shape): diff --git a/src/maxtext/checkpoint_conversion/utils/utils.py b/src/maxtext/checkpoint_conversion/utils/utils.py index 75ee3dffcc..2a7d84c3a7 100644 --- a/src/maxtext/checkpoint_conversion/utils/utils.py +++ b/src/maxtext/checkpoint_conversion/utils/utils.py @@ -1190,6 +1190,7 @@ def save_weights_to_checkpoint( device_count: int, use_ocdbt: bool, use_zarr3: bool, + config=None, ): """Saves model weights to a MaxText-compatible checkpoint with optional sharding. @@ -1205,6 +1206,7 @@ def save_weights_to_checkpoint( use_ocdbt: If True, enables the Optimized Checkpoint Database with Transactions (OCDBT) format for improved metadata handling. use_zarr3: If True, uses the Zarr3 storage format for the underlying array data. + config: Optional config to save along with checkpoint metadata. """ mem_info = psutil.Process() logging.debug("Memory usage: %f GB", mem_info.memory_info().rss / (1024**3)) @@ -1241,7 +1243,7 @@ def save_weights_to_checkpoint( ) logging.debug("Memory usage: %f GB", mem_info.memory_info().rss / (1024**3)) - if checkpointing.save_checkpoint(checkpoint_manager, step_number_to_save_new_ckpt, state_new): + if checkpointing.save_checkpoint(checkpoint_manager, step_number_to_save_new_ckpt, state_new, config=config): max_logging.log(f"saved a checkpoint at step {step_number_to_save_new_ckpt}") # Upon preemption, exit when and only when all ongoing saves are complete. checkpoint_manager.wait_until_finished() diff --git a/src/maxtext/common/checkpointing.py b/src/maxtext/common/checkpointing.py index c778f92bf9..0ac9e7a739 100644 --- a/src/maxtext/common/checkpointing.py +++ b/src/maxtext/common/checkpointing.py @@ -19,7 +19,6 @@ from typing import Any, Optional from absl import flags -import contextlib import datetime from etils import epath from flax import nnx @@ -245,6 +244,27 @@ def _load_linen_checkpoint_into_nnx( return _populate_pure_dict_from_partial(nnx_abstract_pure, partial_nnx) +def _restore_emergency_linen_checkpoint_into_nnx( + checkpoint_manager, + step, + abstract_nnx_state, + map_to_pspec, +): + """Restores an emergency Linen-layout checkpoint into an NNX state.""" + max_logging.log(f"Restoring emergency Linen-layout checkpoint into NNX state at step {step}") + nnx_abstract_pure = abstract_nnx_state.to_pure_dict() + linen_abstract = train_state_nnx.to_linen_checkpoint_dict(nnx_abstract_pure) + restore_args = jax.tree_util.tree_map(map_to_pspec, linen_abstract) + checkpoint_args = ocp.args.PyTreeRestore( + item=linen_abstract, + restore_args=restore_args, + partial_restore=True, + ) + restored = checkpoint_manager.restore(step, args=Composite(state=checkpoint_args)).state + partial_nnx = train_state_nnx.from_linen_checkpoint_dict(restored) + return _populate_pure_dict_from_partial(nnx_abstract_pure, partial_nnx) + + def _rebuild_nnx_with_values(abstract_nnx_state, concrete_weights): """Fills each Variable in `abstract_nnx_state` with the matching restored array.""" leaves, treedef = jax.tree_util.tree_flatten(abstract_nnx_state, is_leaf=lambda x: isinstance(x, nnx.Variable)) @@ -640,14 +660,12 @@ def _restore_grain_iterator( if isinstance(data_iterator, RemoteIteratorWrapper): grain_restore_args = GrainCheckpointRestore(item=data_iterator) restored_state = checkpoint_manager.restore(step, args=Composite(items=checkpoint_args, iter=grain_restore_args)) - _assert_no_shaped_dtype_struct(restored_state) return (restored_state, None) # ElasticIterator: one shared `process_0.json` regardless of shard count. if not isinstance(data_iterator, list) and isinstance(data_iterator.local_iterator, ElasticIterator): grain_restore_args = GrainCheckpointRestore(item=data_iterator.local_iterator) restored_state = checkpoint_manager.restore(step, args=Composite(items=checkpoint_args, iter=grain_restore_args)) - _assert_no_shaped_dtype_struct(restored_state) return (restored_state, None) directory = checkpoint_manager.directory / str(step) / "iter" @@ -696,67 +714,9 @@ def _restore_grain_iterator( # Call restore once with the composed arguments restored_state = checkpoint_manager.restore(step, args=Composite(items=checkpoint_args, iter=grain_restore_args)) - _assert_no_shaped_dtype_struct(restored_state) return (restored_state, None) -def is_structural_or_shape_mismatch(e: Exception) -> bool: - """Helper to check if an exception is likely a PyTree structure or shape mismatch.""" - if not isinstance(e, (ValueError, TypeError)): - return False - msg = str(e).lower() - mismatch_keywords = [ - "mismatch", - "structure", - "shape", - "tree", - "leaf", - "leaves", - "paths matched", - "shapedtypestruct", - "invalid type", - ] - return any(kw in msg for kw in mismatch_keywords) - - -def _assert_no_shaped_dtype_struct(pytree): - """Asserts that there are no jax.ShapeDtypeStruct leaves in the restored pytree.""" - if isinstance(pytree, jax.ShapeDtypeStruct): - raise ValueError( - "Some parameters in the restored state remained as ShapeDtypeStruct" - f" (indicating structure mismatch): {pytree}." - ) - - if hasattr(pytree, "keys") and hasattr(pytree, "__getitem__"): - for k in pytree.keys(): - _assert_no_shaped_dtype_struct(pytree[k]) - elif isinstance(pytree, (list, tuple)): - for v in pytree: - _assert_no_shaped_dtype_struct(v) - else: - leaves = jax.tree_util.tree_leaves(pytree) - if len(leaves) == 1 and leaves[0] is pytree: - return - for leaf in leaves: - _assert_no_shaped_dtype_struct(leaf) - - -@contextlib.contextmanager -def handle_checkpoint_mismatch(context_name: str, path: str): - """Context manager to intercept PyTree/shape mismatches and raise descriptive errors.""" - try: - yield - except Exception as e: - if is_structural_or_shape_mismatch(e): - raise ValueError( - f"Failed to {context_name} from {path}. This is often caused by a" - " mismatch in the 'scan_layers' configuration (stacked vs unstacked)" - " between your current execution command and the saved checkpoint." - f" Original error: {e}" - ) from e - raise - - def load_state_if_possible( checkpoint_manager: CheckpointManager | None, data_iterator: MultiHostDataLoadIterator | list[MultiHostDataLoadIterator] | None, @@ -839,17 +799,31 @@ def map_to_pspec(data): (EmergencyCheckpointManager, EmergencyReplicatorCheckpointManager), ): checkpoint_path = str(checkpoint_manager.directory / str(step) / "items") - with handle_checkpoint_mismatch("restore NNX checkpoint", checkpoint_path): - restored_nnx = _load_linen_checkpoint_into_nnx( - checkpoint_path, - abstract_unboxed_pre_state, - checkpoint_storage_concurrent_gb, - use_ocdbt, - use_zarr3, - ) - _assert_no_shaped_dtype_struct(restored_nnx) + restored_nnx = _load_linen_checkpoint_into_nnx( + checkpoint_path, + abstract_unboxed_pre_state, + checkpoint_storage_concurrent_gb, + use_ocdbt, + use_zarr3, + ) return ({"items": restored_nnx}, None) + if isinstance(abstract_unboxed_pre_state, nnx.State) and isinstance( + checkpoint_manager, + (EmergencyCheckpointManager, EmergencyReplicatorCheckpointManager), + ): + checkpoint_path = str(checkpoint_manager.directory / str(step)) + restored = _restore_emergency_linen_checkpoint_into_nnx( + checkpoint_manager, + step, + abstract_unboxed_pre_state, + map_to_pspec, + ) + return ( + restored, + None, + ) + # Convert nnx.State to pure dict to match how checkpoints are saved for NNX restore_target = abstract_unboxed_pre_state if isinstance(abstract_unboxed_pre_state, nnx.State): @@ -863,49 +837,46 @@ def map_to_pspec(data): ) checkpoint_path = str(checkpoint_manager.directory / str(step)) - with handle_checkpoint_mismatch("restore checkpoint", checkpoint_path): - match (checkpoint_manager, dataset_type, data_iterator): - # Case 1: Matches if 'checkpoint_manager' is an instance of either EmergencyCheckpointManager - # or EmergencyReplicatorCheckpointManager. The '_' indicates that 'dataset_type' and - # 'data_iterator' can be any value and aren't used in this pattern. - case (checkpoint_manager, _, _) if isinstance( - checkpoint_manager, - ( - EmergencyCheckpointManager, - EmergencyReplicatorCheckpointManager, - ), - ): - restored = checkpoint_manager.restore(step, args=Composite(state=checkpoint_args)).state - _assert_no_shaped_dtype_struct(restored) - return ( - restored, - None, - ) - # Case 2: Matches if dataset type is "grain" and the data iterator is not a - # PlaceHolderDataIterator and a specific checkpoint file exists for the iterator - case ( + match (checkpoint_manager, dataset_type, data_iterator): + # Case 1: Matches if 'checkpoint_manager' is an instance of either EmergencyCheckpointManager + # or EmergencyReplicatorCheckpointManager. The '_' indicates that 'dataset_type' and + # 'data_iterator' can be any value and aren't used in this pattern. + case (checkpoint_manager, _, _) if isinstance( + checkpoint_manager, + ( + EmergencyCheckpointManager, + EmergencyReplicatorCheckpointManager, + ), + ): + restored = checkpoint_manager.restore(step, args=Composite(state=checkpoint_args)).state + return ( + restored, + None, + ) + # Case 2: Matches if dataset type is "grain" and the data iterator is not a + # PlaceHolderDataIterator and a specific checkpoint file exists for the iterator + case ( + checkpoint_manager, + dataset_type, + data_iterator, + ) if ( + dataset_type == "grain" + and data_iterator + and not isinstance(data_iterator, PlaceHolderDataIterator) + and (checkpoint_manager.directory / str(step) / "iter").exists() + ): + return _restore_grain_iterator( checkpoint_manager, - dataset_type, + step, data_iterator, - ) if ( - dataset_type == "grain" - and data_iterator - and not isinstance(data_iterator, PlaceHolderDataIterator) - and (checkpoint_manager.directory / str(step) / "iter").exists() - ): - return _restore_grain_iterator( - checkpoint_manager, - step, - data_iterator, - checkpoint_args, - expansion_factor_real_data, - ) - # Case 3: Default/Fallback case. - # This case acts as a wildcard ('_') and matches if none of the preceding cases were met. - case _: - restored = checkpoint_manager.restore(step, args=Composite(items=checkpoint_args)) - _assert_no_shaped_dtype_struct(restored) - return (restored, None) + checkpoint_args, + expansion_factor_real_data, + ) + # Case 3: Default/Fallback case. + # This case acts as a wildcard ('_') and matches if none of the preceding cases were met. + case _: + restored = checkpoint_manager.restore(step, args=Composite(items=checkpoint_args)) + return (restored, None) if source_checkpoint_layout == "safetensors_dynamic": path = load_parameters_from_path or load_full_state_from_path @@ -918,30 +889,26 @@ def map_to_pspec(data): else: params = abstract_unboxed_pre_state.params - with handle_checkpoint_mismatch("load parameters", load_parameters_from_path): - restored_params = load_params_from_path( - load_parameters_from_path, - params, - checkpoint_storage_concurrent_gb, - use_ocdbt=use_ocdbt, - use_zarr3=use_zarr3, - ) - _assert_no_shaped_dtype_struct(restored_params) + restored_params = load_params_from_path( + load_parameters_from_path, + params, + checkpoint_storage_concurrent_gb, + use_ocdbt=use_ocdbt, + use_zarr3=use_zarr3, + ) return None, restored_params elif load_full_state_from_path != "": max_logging.log(f"Loading full state from path: {load_full_state_from_path}") - with handle_checkpoint_mismatch("load full state", load_full_state_from_path): - restored_state = _load_full_state_from_path( - path=load_full_state_from_path, - abstract_unboxed_pre_state=abstract_unboxed_pre_state, - enable_orbax_v1=enable_orbax_v1, - checkpoint_conversion_fn=checkpoint_conversion_fn, - source_checkpoint_layout=source_checkpoint_layout, - checkpoint_storage_concurrent_gb=checkpoint_storage_concurrent_gb, - use_ocdbt=use_ocdbt, - use_zarr3=use_zarr3, - ) - _assert_no_shaped_dtype_struct(restored_state) + restored_state = _load_full_state_from_path( + path=load_full_state_from_path, + abstract_unboxed_pre_state=abstract_unboxed_pre_state, + enable_orbax_v1=enable_orbax_v1, + checkpoint_conversion_fn=checkpoint_conversion_fn, + source_checkpoint_layout=source_checkpoint_layout, + checkpoint_storage_concurrent_gb=checkpoint_storage_concurrent_gb, + use_ocdbt=use_ocdbt, + use_zarr3=use_zarr3, + ) return {"items": restored_state}, None else: max_logging.log("No existing checkpoints found, not restoring checkpoint.") @@ -1023,6 +990,52 @@ def save_params_to_path(checkpoint_dir, params, use_ocdbt=True, use_zarr3=True): print(f"Quantized params checkpoint saved at: {checkpoint_dir}") +def load_checkpoint_metadata(checkpoint_dir_path: str) -> dict[str, Any]: + """Loads custom metadata from an Orbax checkpoint. + + Args: + checkpoint_dir_path: Path to the checkpoint directory. + + Returns: + A dictionary containing custom metadata, or an empty dictionary if none is + present or loading fails. + """ + checkpoint_dir = epath.Path(checkpoint_dir_path) + try: + ckptr = ocp.StandardCheckpointer() + metadata = ckptr.metadata(checkpoint_dir) + return metadata.custom_metadata or {} + except Exception as e: # pylint: disable=broad-except + max_logging.log(f"Warning: Failed to load checkpoint metadata: {e}") + return {} + + +def _uses_local_checkpoint_period(config): + return config.enable_emergency_checkpoint or config.enable_multi_tier_checkpointing + + +def _should_save_checkpoint_at_step(checkpoint_manager, step, config, force): + """Returns whether MaxText should build and dispatch checkpoint args.""" + if force: + return True + if config.enable_continuous_checkpointing: + base_checkpoint_due = bool(checkpoint_manager.should_save(step)) + else: + base_checkpoint_due = step % config.checkpoint_period == 0 + local_checkpoint_due = _uses_local_checkpoint_period(config) and step % config.local_checkpoint_period == 0 + autocheckpoint_due = config.enable_autocheckpoint and checkpoint_manager.reached_preemption(step) + return base_checkpoint_due or local_checkpoint_due or autocheckpoint_due + + +def _handle_post_checkpoint_preemption(checkpoint_manager, step, force_ckpt_save): + """Waits on final/preemption saves and raises if preempted.""" + reached_preemption = checkpoint_manager.reached_preemption(step) + if force_ckpt_save or reached_preemption: + checkpoint_manager.wait_until_finished() + if reached_preemption: + raise exceptions.StopTraining("Job is preempted.") + + def maybe_save_checkpoint(checkpoint_manager, state, config, data_iterator, step=None): """Save checkpoint if checkpointing is enabled.""" if checkpoint_manager is None: @@ -1041,6 +1054,18 @@ def maybe_save_checkpoint(checkpoint_manager, state, config, data_iterator, step # Linen TrainState has .step attribute actual_step = int(state.step) - 1 + # Determine if a checkpoint save should be forced, overriding the usual + # `config.checkpoint_period` logic. + # This occurs if this function was called: + # without an explicit 'step' (implying it's a checkpoint save for final step), + # AND the 'actual_step' is a valid step, + # AND it's not a step that would normally trigger a checkpoint save. + force_ckpt_save = step is None and actual_step != -1 and (actual_step % config.checkpoint_period != 0) + + if not _should_save_checkpoint_at_step(checkpoint_manager, actual_step, config, force_ckpt_save): + _handle_post_checkpoint_preemption(checkpoint_manager, actual_step, force_ckpt_save) + return + if checkpoint_manager.latest_step() == actual_step: max_logging.log(f"Checkpoint for step {actual_step} already exists, skipping save.") return @@ -1055,13 +1080,6 @@ def maybe_save_checkpoint(checkpoint_manager, state, config, data_iterator, step else: state = train_state_nnx.to_linen_checkpoint_dict(state.to_pure_dict()) - # Determine if a checkpoint save should be forced, overriding the usual `config.checkpoint_period` logic. - # This occurs if this function was called: - # without an explicit 'step' (implying it's a checkpoint save for final step), - # AND the 'actual_step' is a valid step, - # AND it's not a step that would normally trigger a checkpoint save. - force_ckpt_save = step is None and actual_step != -1 and (actual_step % config.checkpoint_period != 0) - try: checkpoint_saved = save_checkpoint(checkpoint_manager, actual_step, state, config, data_iterator, force_ckpt_save) if checkpoint_saved: @@ -1083,13 +1101,9 @@ def maybe_save_checkpoint(checkpoint_manager, state, config, data_iterator, step except Exception as e: raise exceptions.StopTraining(f"Checkpointing failed. {str(e)}") from e - # Wait for any pending checkpoint save to finish during preemption or final step save - if force_ckpt_save or checkpoint_manager.reached_preemption(actual_step): - checkpoint_manager.wait_until_finished() - - # Raise exception upon preemption - if checkpoint_manager.reached_preemption(actual_step): - raise exceptions.StopTraining("Job is preempted.") + # Wait for any pending checkpoint save to finish during preemption or final + # step save, then raise upon preemption. + _handle_post_checkpoint_preemption(checkpoint_manager, actual_step, force_ckpt_save) def save_checkpoint(checkpoint_manager, step, state, config=None, data_iterator=None, force=False): @@ -1098,7 +1112,7 @@ def save_checkpoint(checkpoint_manager, step, state, config=None, data_iterator= if ( force or (step % config.checkpoint_period == 0 and not config.enable_continuous_checkpointing) - or (config.enable_emergency_checkpoint and step % config.local_checkpoint_period == 0) + or (_uses_local_checkpoint_period(config) and step % config.local_checkpoint_period == 0) or (config.enable_autocheckpoint and checkpoint_manager.reached_preemption(step)) ): blocking_until_ready_start = time.time() @@ -1142,6 +1156,13 @@ def save_checkpoint(checkpoint_manager, step, state, config=None, data_iterator= grain_iters_to_save.append((data_iter.local_iterator, process_index, process_count_total)) save_args_composite["iter"] = GrainCheckpointSave(item=grain_iters_to_save) + custom_metadata = {} + if config: + if hasattr(config, "scan_layers"): + custom_metadata["scan_layers"] = config.scan_layers + if hasattr(config, "lora") and config.lora and getattr(config.lora, "lora_rank", 0) > 0: + custom_metadata["lora"] = config.lora.model_dump() + match (checkpoint_manager, config, data_iterator): case (checkpoint_manager, _, _) if isinstance( checkpoint_manager, (EmergencyCheckpointManager, EmergencyReplicatorCheckpointManager) @@ -1149,4 +1170,6 @@ def save_checkpoint(checkpoint_manager, step, state, config=None, data_iterator= replicator_error_handler(config) return checkpoint_manager.save(step, args=Composite(state=checkpoint_args), force=force) case _: - return checkpoint_manager.save(step, args=Composite(**save_args_composite), force=force) + return checkpoint_manager.save( + step, args=Composite(**save_args_composite), force=force, custom_metadata=custom_metadata + ) diff --git a/src/maxtext/common/data_loader.py b/src/maxtext/common/data_loader.py index 21bd870bc8..0cfa3069e3 100644 --- a/src/maxtext/common/data_loader.py +++ b/src/maxtext/common/data_loader.py @@ -26,6 +26,7 @@ from maxtext.trainers.diloco import diloco from maxtext.utils import exceptions from maxtext.utils.sharding import get_input_data_sharding +from maxtext.src.maxtext.utils import elastic_utils class DataLoader: @@ -63,6 +64,11 @@ def load_next_batch_pre_sharding(self): self.last_batch = example_batch self.check_example_batch() except Exception as e: # pylint: disable=broad-except + if elastic_utils.elastic_enabled(self.config) and ( + isinstance(e, jax.errors.JaxRuntimeError) + or isinstance(e, elastic_utils.manager.ScaleUpSignalError) + ): + raise if isinstance(e, StopIteration): raise exceptions.StopTraining(f"You may have run out of training data. Received {type(e)} exception: ({e})") else: diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml index 82e448c530..aaa7c1d8fe 100644 --- a/src/maxtext/configs/base.yml +++ b/src/maxtext/configs/base.yml @@ -159,6 +159,8 @@ weight_sparsity_m: null weight_sparsity_update_step: 10 # The first number of steps before updating the sparsity masks. weight_sparsity_start_step: 50 +# Whether to use Transformer Engine's collective GEMM overlap algorithm with quantization in MLP layers. Note: only valid with tensor-sequence parallelism. +use_te_comm_gemm_overlap: false decoder_block: "llama2" # which style of decoderblock to use. # global parameter scale needs to be a power of 2. if you want finer grained control of the model sizes @@ -365,7 +367,8 @@ engram: 'remat' optimizer_memory_host_offload: false parameter_memory_host_offload: false -scan_layers: true # We recommend setting this to false when using pipeline parallelism, instead scanning the PP iterations. +# Whether to use jax.lax.scan over layers (stacked/unstacked checkpoint). We recommend setting this to false when using pipeline parallelism, instead scanning the PP iterations. When resuming from a checkpoint, this flag is auto-determined from metadata. +scan_layers: true param_scan_axis: 1 # The attention parameter dictates the specific algorithm/methodology used to compute the attention scores @@ -492,7 +495,7 @@ compile_xla_flags: "" # Compiler options e.g. compile_xla_flags="--xla_tpu_num_s # Parallelism shard_mode: "auto" # can be either auto or explicit custom_mesh_and_rule: "" # replace default mesh and logical rule by specifying yml name under config/mesh_and_rule/. -mesh_axes: ['diloco', 'data', 'stage', 'fsdp', 'fsdp_transpose', 'context', 'context_autoregressive', 'tensor', 'tensor_transpose', 'tensor_sequence', 'expert', 'autoregressive'] +mesh_axes: ['diloco', 'data', 'stage', 'fsdp', 'fsdp_transpose', 'context', 'context_autoregressive', 'tensor', 'tensor_sequence', 'expert', 'autoregressive'] logical_axis_rules: [ ['circular_repeats', []], # ========================================== @@ -501,40 +504,36 @@ logical_axis_rules: [ # Vocab Activations ['activation_embed_and_logits_batch', ['data', 'stage', 'fsdp', 'fsdp_transpose', 'expert']], ['activation_embed_and_logits_batch_sequence', ['data', 'stage', 'fsdp', 'fsdp_transpose', 'context', 'expert']], - ['activation_vocab', ['tensor', 'tensor_transpose', 'tensor_sequence']], - ['activation_vocab', ['tensor', 'tensor_transpose']], + ['activation_vocab', ['tensor', 'tensor_sequence']], + ['activation_vocab', ['tensor']], ['activation_vocab', 'tensor_sequence'], # Vocab Weights - ['vocab', ['tensor', 'tensor_transpose', 'tensor_sequence', 'autoregressive']], + ['vocab', ['tensor', 'tensor_sequence', 'autoregressive']], ['embed_vocab', ['fsdp', 'fsdp_transpose', 'context', 'expert']], # ========================================== # Attention # ========================================== # Attention Activations ['activation_batch_attn', ['data', 'fsdp', 'fsdp_transpose', 'expert']], - ['activation_heads', ['tensor', 'tensor_transpose', 'tensor_sequence', 'autoregressive']], - ['activation_kv_heads', ['tensor', 'tensor_transpose', 'tensor_sequence']], + ['activation_heads', ['tensor', 'tensor_sequence', 'autoregressive']], + ['activation_kv_heads', ['tensor', 'tensor_sequence']], ['activation_length_attn', ['context']], ['activation_q_length', ['context']], ['activation_kv_length', []], - ['activation_embed_attn', ['tensor', 'tensor_transpose']], - ['activation_kv', ['tensor', 'tensor_transpose', 'tensor_sequence']], + ['activation_embed_attn', ['tensor']], + ['activation_kv', ['tensor', 'tensor_sequence']], ['activation_kv_batch', ['data', 'fsdp', 'fsdp_transpose', 'expert']], - ['activation_kv_head_dim', ['tensor', 'tensor_transpose', 'tensor_sequence']], + ['activation_kv_head_dim', ['tensor', 'tensor_sequence']], # Attention Weights - ['heads', ['tensor', 'tensor_transpose', 'tensor_sequence', 'autoregressive']], - ['q_heads', ['tensor', 'tensor_transpose', 'tensor_sequence', 'autoregressive']], - ['kv_heads', ['tensor', 'tensor_transpose', 'tensor_sequence', 'autoregressive']], + ['heads', ['tensor', 'tensor_sequence', 'autoregressive']], + ['q_heads', ['tensor', 'tensor_sequence', 'autoregressive']], + ['kv_heads', ['tensor', 'tensor_sequence', 'autoregressive']], ['qkv', []], ['kv', []], ['kv_head_dim', []], - ['q_lora', ['fsdp', 'fsdp_transpose', 'context', 'tensor_transpose', 'expert']], - ['q_lora', ['fsdp', 'context', 'tensor_transpose', 'expert']], ['q_lora', ['fsdp', 'fsdp_transpose', 'context', 'expert']], ['q_lora', ['fsdp', 'context', 'expert']], ["q_lora_up_proj", []], - ['kv_lora', ['fsdp', 'fsdp_transpose', 'context', 'tensor_transpose', 'expert']], - ['kv_lora', ['fsdp', 'context', 'tensor_transpose', 'expert']], ['kv_lora', ['fsdp', 'fsdp_transpose', 'context', 'expert']], ['kv_lora', ['fsdp', 'context', 'expert']], ["kv_lora_up_proj", []], @@ -545,37 +544,33 @@ logical_axis_rules: [ ['activation_batch_moe', ['data', 'fsdp', 'fsdp_transpose', 'expert']], ['activation_length_moe', ['context']], ['activation_norm_length_moe', ['tensor_sequence', 'context']], - ['activation_embed_moe', ['tensor', 'tensor_transpose']], - ['activation_mlp_moe', ['tensor', 'tensor_transpose', 'tensor_sequence']], + ['activation_embed_moe', ['tensor']], + ['activation_mlp_moe', ['tensor', 'tensor_sequence']], ['activation_exp', ['expert']], # MoE Weights ['exp', 'expert'], ['mlp_moe', ['fsdp_transpose', 'tensor', 'tensor_sequence', 'autoregressive']], - ['embed_moe', ['fsdp', 'fsdp_transpose', 'tensor_transpose', 'context']], - ['embed_moe', ['fsdp', 'tensor_transpose', 'context']], ['embed_moe', ['fsdp', 'fsdp_transpose', 'context']], ['embed_moe', ['fsdp', 'context']], # ========================================== # Standard MLP / Dense Layers / Model Structure # ========================================== # Dense Activations - ['activation_mlp', ['tensor', 'tensor_transpose', 'tensor_sequence']], + ['activation_mlp', ['tensor', 'tensor_sequence']], # Note activation batch and length also get used in vocab ['activation_batch', ['data', 'fsdp', 'fsdp_transpose', 'expert']], ['activation_length', ['context']], ['activation_norm_length', ['tensor_sequence', 'context']], - ['activation_embed', ['tensor', 'tensor_transpose']], + ['activation_embed', ['tensor']], ['activation_stage', 'stage'], # General Weights ['mlp', ['fsdp_transpose', 'tensor', 'tensor_sequence', 'autoregressive']], # GDN (linear-attention) projections shard like 'mlp' during training; the # vLLM serving config overrides this to match tpu-inference's ATTN_HEAD order. ['gdn_head', ['fsdp_transpose', 'tensor', 'tensor_sequence', 'autoregressive']], - ['embed', ['fsdp', 'fsdp_transpose', 'tensor_transpose', 'context', 'expert']], - ['embed', ['fsdp', 'tensor_transpose', 'context', 'expert']], ['embed', ['fsdp', 'fsdp_transpose', 'context', 'expert']], ['embed', ['fsdp', 'context', 'expert']], - ['norm', ['tensor', 'tensor_transpose']], + ['norm', ['tensor']], ['layers', 'stage'], ['diloco', 'diloco'], ['engram_dim', ['tensor']], @@ -590,7 +585,6 @@ logical_axis_rules: [ ['activation_prefill_kv_batch', ['data', 'fsdp', 'fsdp_transpose', 'expert']], ['decode_batch', ['data', 'fsdp', 'fsdp_transpose', 'expert']], ['decode_length', []], - ['cache_heads', ['autoregressive', 'tensor', 'tensor_transpose', 'tensor_sequence']], ['cache_heads', ['autoregressive', 'tensor', 'tensor_sequence']], ['paged_kv_heads', ['tensor']], ['cache_batch_prefill', []], @@ -605,11 +599,10 @@ logical_axis_rules: [ # Deprecated / Scheduled for Removal # ========================================== ['mlp_no_fsdp', ['tensor', 'tensor_sequence', 'autoregressive']], - ['embed_tensor_transpose', ['tensor_transpose']], ['exp_with_fsdp', 'fsdp'], ] # Axes used for DCN must be earlier in this list than ICI, see (b/339009148) for details -data_sharding: [['data', 'stage', 'fsdp', 'fsdp_transpose', 'context', 'context_autoregressive', 'tensor', 'tensor_transpose', 'tensor_sequence', 'expert', 'autoregressive']] +data_sharding: [['data', 'stage', 'fsdp', 'fsdp_transpose', 'context', 'context_autoregressive', 'tensor', 'tensor_sequence', 'expert', 'autoregressive']] input_data_sharding_logical_axes: ['activation_embed_and_logits_batch', 'activation_norm_length'] # Determines which physical axis plays the role of context parallelism for input data processing and load balancing # only supports "context" or "expert" (when custom_mesh_and_rule=ep-as-cp) @@ -630,7 +623,6 @@ dcn_sequence_parallelism: 1 # never recommended dcn_context_parallelism: 1 dcn_context_autoregressive_parallelism: 1 dcn_tensor_parallelism: 1 # never recommended -dcn_tensor_transpose_parallelism: 1 dcn_tensor_sequence_parallelism: 1 # never recommended dcn_pipeline_parallelism: 1 dcn_expert_parallelism: 1 @@ -643,7 +635,6 @@ ici_sequence_parallelism: 1 ici_context_parallelism: 1 ici_context_autoregressive_parallelism: 1 ici_tensor_parallelism: 1 -ici_tensor_transpose_parallelism: 1 ici_tensor_sequence_parallelism: 1 ici_autoregressive_parallelism: 1 ici_pipeline_parallelism: 1 @@ -890,7 +881,7 @@ diloco_sync_period: 36 diloco_outer_lr: 0.3 diloco_outer_momentum: 0.9 # DCN bandwidth throttling parameters (used for simulating slow networks). -# If dcn_bandwidth_limit is empty, no throttling is applied. +# If dcn_bandwidth_limit is empty, no throttling is applied. **Per-VM** throughput limit. Example: 10gbit. dcn_bandwidth_limit: "" # The burst size parameter for the traffic control (tc) token bucket filter. dcn_bandwidth_burst: "10mb" diff --git a/src/maxtext/configs/decoupled_base_test.yml b/src/maxtext/configs/decoupled_base_test.yml index 8494011ccc..485bb4e8ae 100644 --- a/src/maxtext/configs/decoupled_base_test.yml +++ b/src/maxtext/configs/decoupled_base_test.yml @@ -42,7 +42,6 @@ ici_pipeline_parallelism: 1 ici_expert_parallelism: 1 ici_sequence_parallelism: 1 ici_context_parallelism: 1 -ici_tensor_transpose_parallelism: 1 ici_tensor_sequence_parallelism: 1 ici_autoregressive_parallelism: 1 ici_fsdp_parallelism: -1 @@ -57,7 +56,6 @@ dcn_pipeline_parallelism: 1 dcn_expert_parallelism: 1 dcn_sequence_parallelism: 1 dcn_context_parallelism: 1 -dcn_tensor_transpose_parallelism: 1 dcn_tensor_sequence_parallelism: 1 dcn_autoregressive_parallelism: 1 dcn_fsdp_parallelism: 1 diff --git a/src/maxtext/configs/inference/vllm.yml b/src/maxtext/configs/inference/vllm.yml index 157845c5b7..44efdcb161 100644 --- a/src/maxtext/configs/inference/vllm.yml +++ b/src/maxtext/configs/inference/vllm.yml @@ -88,6 +88,9 @@ logical_axis_rules: [ ['mlp', ['attn_dp', 'model']], ['embed', []], ['norm', []], + ['layers', []], + ['dense_layers', []], + ['moe_layers', []], # ========================================== # Inference(Prefill, Decode, Cache) # ========================================== diff --git a/src/maxtext/configs/post_train/rl.yml b/src/maxtext/configs/post_train/rl.yml index 3ecdd3b981..64460ea066 100644 --- a/src/maxtext/configs/post_train/rl.yml +++ b/src/maxtext/configs/post_train/rl.yml @@ -66,8 +66,6 @@ rl: off_policy_steps: 0 # System prompt injected into the agent at rollout time. system_prompt: '' - # If true, mask degenerate groups (all-zero advantages) from contributing to the loss. - degenerate_group_masking: true # Upper-bound clipping epsilon for GRPO loss; defaults to grpo_epsilon when null. epsilon_high: null # Number of model keys to chunk for resharding tensors between trainer and rollout devices. @@ -237,6 +235,13 @@ eval_dataset_name: 'openai/gsm8k' # tmvp_config, x) -> dict with keys {prompts, question, answer}. When empty # (default), the built-in utils_rl.process_data is used. dataset_processor_path: '' +# Optional: path to a user-provided Python file with custom reward functions, +# AND a comma-separated list of function names to import from that file. +# Each function signature: fn(prompts, completions, tmvp_config, **kwargs) -> list[float]. +# When both are set, the built-in [match_format_exactly, match_format_approximately, +# check_numbers] reward list is REPLACED entirely. +reward_functions_path: '' +reward_functions: '' train_split: 'train' eval_split: 'test' hf_name: 'main' # subset of Hugging Face dataset diff --git a/src/maxtext/configs/pyconfig.py b/src/maxtext/configs/pyconfig.py index f4a8e3965f..a23c3e3eb8 100644 --- a/src/maxtext/configs/pyconfig.py +++ b/src/maxtext/configs/pyconfig.py @@ -498,6 +498,7 @@ def _initialize_pydantic(argv: list[str] | None = None, **kwargs) -> MaxTextConf # 6. Handle environment variable overrides cli_keys = frozenset(omegaconf.OmegaConf.to_container(cli_cfg, resolve=True).keys()) kwargs_keys = frozenset(kwargs.keys()) + env_keys = set() for k in tuple(raw_keys_dict.keys()): env_key = yaml_key_to_env_key(k) if env_key in os.environ: @@ -513,6 +514,7 @@ def _initialize_pydantic(argv: list[str] | None = None, **kwargs) -> MaxTextConf "This is not allowed, unless setting `override_model_config=True`." ) + env_keys.add(k) new_proposal = os.environ.get(env_key) original_value = raw_keys_dict.get(k) parser = None @@ -551,6 +553,8 @@ def _initialize_pydantic(argv: list[str] | None = None, **kwargs) -> MaxTextConf compilation_cache.set_cache_dir(os.path.expanduser(pydantic_kwargs["jax_cache_dir"])) pydantic_config = types.MaxTextConfig(**pydantic_kwargs) + explicit_keys = set((cli_keys | kwargs_keys | env_keys) & frozenset(types.MaxTextConfig.model_fields.keys())) + object.__setattr__(pydantic_config, "__pydantic_fields_set__", explicit_keys) config = HyperParameters(pydantic_config) if config.log_config: diff --git a/src/maxtext/configs/pyconfig_deprecated.py b/src/maxtext/configs/pyconfig_deprecated.py index 3e6329b7aa..ffced43785 100644 --- a/src/maxtext/configs/pyconfig_deprecated.py +++ b/src/maxtext/configs/pyconfig_deprecated.py @@ -249,12 +249,11 @@ def validate_keys(keys): keys["global_rampup_samples"], ) - context_parallel_size = get_context_parallel_size(keys) # TODO remove after b/435512699 resolved - if context_parallel_size > 1 and keys["context_parallel_load_balance"] and keys["attention_type"] == "chunk": - raise ValueError("Currently load-balanced context parallelism is not supported for chunk attention.") - validate_context_parallel_strategy_ring(context_parallel_size, keys["context_parallel_strategy"], keys["hardware"]) + validate_context_parallel_strategy_ring( + keys["context_parallel_size"], keys["context_parallel_strategy"], keys["hardware"] + ) if keys["mtp_eval_target_module"] < 0: raise ValueError("mtp_eval_target_module cannot be negative. Set to 0 to disable evaluation.") @@ -816,6 +815,7 @@ def user_init(raw_keys): raw_keys["num_slices"] = max_utils.get_num_slices(raw_keys) raw_keys["quantization_local_shard_count"] = get_quantization_local_shard_count(raw_keys) + raw_keys["context_parallel_size"] = get_context_parallel_size(raw_keys) raw_keys = create_parallelisms_list(raw_keys) raw_keys = set_and_validate_pipeline_config(raw_keys) @@ -918,7 +918,6 @@ def create_parallelisms_list(raw_keys): raw_keys["ici_context_parallelism"], raw_keys["ici_context_autoregressive_parallelism"], raw_keys["ici_tensor_parallelism"], - raw_keys["ici_tensor_transpose_parallelism"], raw_keys["ici_tensor_sequence_parallelism"], raw_keys["ici_expert_parallelism"], raw_keys["ici_autoregressive_parallelism"], @@ -932,7 +931,6 @@ def create_parallelisms_list(raw_keys): raw_keys["dcn_context_parallelism"], raw_keys["dcn_context_autoregressive_parallelism"], raw_keys["dcn_tensor_parallelism"], - raw_keys["dcn_tensor_transpose_parallelism"], raw_keys["dcn_tensor_sequence_parallelism"], raw_keys["dcn_expert_parallelism"], raw_keys["dcn_autoregressive_parallelism"], @@ -1029,7 +1027,6 @@ def pipeline_first_axis(raw_keys): raw_keys["ici_context_parallelism"], raw_keys["ici_context_autoregressive_parallelism"], raw_keys["ici_tensor_parallelism"], - raw_keys["ici_tensor_transpose_parallelism"], raw_keys["ici_tensor_sequence_parallelism"], raw_keys["ici_expert_parallelism"], raw_keys["ici_autoregressive_parallelism"], @@ -1043,7 +1040,6 @@ def pipeline_first_axis(raw_keys): raw_keys["dcn_context_parallelism"], raw_keys["dcn_context_autoregressive_parallelism"], raw_keys["dcn_tensor_parallelism"], - raw_keys["dcn_tensor_transpose_parallelism"], raw_keys["dcn_tensor_sequence_parallelism"], raw_keys["dcn_expert_parallelism"], raw_keys["dcn_autoregressive_parallelism"], @@ -1057,7 +1053,6 @@ def pipeline_first_axis(raw_keys): "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive", @@ -1072,7 +1067,6 @@ def pipeline_first_axis(raw_keys): "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive", @@ -1217,9 +1211,7 @@ def validate_shard_expert_on_fsdp(raw_keys): if raw_keys["shard_exp_on_fsdp"] and raw_keys["num_experts"] % raw_keys["ici_fsdp_parallelism"] != 0: raise ValueError("shard_exp_on_fsdp requires num_experts is divisiable by ici_fsdp_parallelism.") if raw_keys["shard_exp_on_fsdp"] and (using_tensor_parallelism(raw_keys) or using_expert_parallelism(raw_keys)): - raise ValueError( - "shard_exp_on_fsdp requires ici_expert_parallelism = 1 and ici_tensor_parallelism/ici_tensor_transpose_parallelism = 1." - ) + raise ValueError("shard_exp_on_fsdp requires ici_expert_parallelism = 1 and ici_tensor_parallelism = 1.") def validate_ragged_dot(raw_keys): diff --git a/src/maxtext/configs/tpu/v4/22b.sh b/src/maxtext/configs/tpu/v4/22b.sh index 35798dbdc1..83fad6c38e 100644 --- a/src/maxtext/configs/tpu/v4/22b.sh +++ b/src/maxtext/configs/tpu/v4/22b.sh @@ -18,11 +18,10 @@ echo "Running 22b.sh" # # Command Flags: # OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml) -# DATASET_PATH (Required, unless dataset_path is already set in base.yml) # RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE) # # Example to invoke this script for training: -# bash src/maxtext/configs/tpu/v4/22b.sh RUN_NAME="" OUTPUT_PATH="gs://" DATASET_PATH="gs://" +# bash src/maxtext/configs/tpu/v4/22b.sh RUN_NAME="" OUTPUT_PATH="gs://" # # Example to AOT compile: # bash src/maxtext/configs/tpu/v4/22b.sh EXECUTABLE=train_compile M_COMPILE_TOPOLOGY=v4-128 M_COMPILE_TOPOLOGY_NUM_SLICES=2 @@ -58,4 +57,4 @@ export LIBTPU_INIT_ARGS="--xla_enable_async_all_gather=true TPU_MEGACORE=MEGACOR python3 -m maxtext.trainers.pre_train.$EXECUTABLE "${MAXTEXT_CONFIGS_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/maxtext/configs}"//base.yml\ ici_fsdp_parallelism=64 steps=10 per_device_batch_size=13 profiler=xplane remat_policy=full attention=flash num_vocab_tiling=8\ base_emb_dim=6144 base_num_kv_heads=24 base_num_query_heads=24 base_mlp_dim=24576 base_num_decoder_layers=48\ - base_output_directory=$OUTPUT_PATH dataset_path=$DATASET_PATH + base_output_directory=$OUTPUT_PATH dataset_type=synthetic diff --git a/src/maxtext/configs/tpu/v4/52b.sh b/src/maxtext/configs/tpu/v4/52b.sh index 702a4812ef..f774679c50 100644 --- a/src/maxtext/configs/tpu/v4/52b.sh +++ b/src/maxtext/configs/tpu/v4/52b.sh @@ -18,11 +18,10 @@ echo "Running 52b.sh" # # Command Flags: # OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml) -# DATASET_PATH (Required, unless dataset_path is already set in base.yml) # RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE) # # Example to invoke this script: -# bash src/maxtext/configs/tpu/v4/52b.sh RUN_NAME="" OUTPUT_PATH="gs://" DATASET_PATH="gs://" +# bash src/maxtext/configs/tpu/v4/52b.sh RUN_NAME="" OUTPUT_PATH="gs://" # # Example to AOT compile: # bash src/maxtext/configs/tpu/v4/52b.sh EXECUTABLE=train_compile M_COMPILE_TOPOLOGY=v4-384 M_COMPILE_TOPOLOGY_NUM_SLICES=2 @@ -59,4 +58,4 @@ python3 -m maxtext.trainers.pre_train.$EXECUTABLE "${MAXTEXT_CONFIGS_DIR:-${MAXT profiler=xplane enable_checkpointing=false steps=10\ ici_fsdp_parallelism=192 ici_tensor_parallelism=1 per_device_batch_size=7 remat_policy=full attention=flash num_vocab_tiling=8\ base_num_decoder_layers=32 base_emb_dim=12288 base_mlp_dim=49152 base_num_query_heads=32 base_num_kv_heads=32 learning_rate=1e-8\ - base_output_directory=$OUTPUT_PATH dataset_path=$DATASET_PATH + base_output_directory=$OUTPUT_PATH dataset_type=synthetic diff --git a/src/maxtext/configs/tpu/v5e/128b.sh b/src/maxtext/configs/tpu/v5e/128b.sh index 069abf376e..2fdc307b25 100644 --- a/src/maxtext/configs/tpu/v5e/128b.sh +++ b/src/maxtext/configs/tpu/v5e/128b.sh @@ -4,11 +4,10 @@ echo "Running 128b.sh" # # Command Flags: # OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml) -# DATASET_PATH (Required, unless dataset_path is already set in base.yml) # RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE) # # Example to invoke this script: -# bash src/maxtext/configs/tpu/v5e/128b.sh RUN_NAME="" OUTPUT_PATH="gs://" DATASET_PATH="gs://" +# bash src/maxtext/configs/tpu/v5e/128b.sh RUN_NAME="" OUTPUT_PATH="gs://" # # Example to AOT compile: # bash src/maxtext/configs/tpu/v5e/128b.sh EXECUTABLE=train_compile M_COMPILE_TOPOLOGY=v5e-256 M_COMPILE_TOPOLOGY_NUM_SLICES=2 diff --git a/src/maxtext/configs/tpu/v5e/16b.sh b/src/maxtext/configs/tpu/v5e/16b.sh index 4c45661146..da66a72adf 100644 --- a/src/maxtext/configs/tpu/v5e/16b.sh +++ b/src/maxtext/configs/tpu/v5e/16b.sh @@ -4,11 +4,10 @@ echo "Running 16b.sh" # # Command Flags: # OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml) -# DATASET_PATH (Required, unless dataset_path is already set in base.yml) # RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE) # # Example to invoke this script: -# bash src/maxtext/configs/tpu/v5e/16b.sh RUN_NAME="" OUTPUT_PATH="gs://" DATASET_PATH="gs://" +# bash src/maxtext/configs/tpu/v5e/16b.sh RUN_NAME="" OUTPUT_PATH="gs://" # # Example to AOT compile: # bash src/maxtext/configs/tpu/v5e/16b.sh EXECUTABLE=train_compile M_COMPILE_TOPOLOGY=v5e-256 M_COMPILE_TOPOLOGY_NUM_SLICES=2 @@ -45,5 +44,5 @@ python3 -m maxtext.trainers.pre_train.$EXECUTABLE "${MAXTEXT_CONFIGS_DIR:-${MAXT steps=15 per_device_batch_size=6 enable_checkpointing=false\ remat_policy=full global_parameter_scale=16\ max_target_length=2048 base_output_directory=$OUTPUT_PATH\ - dataset_path=$DATASET_PATH use_iota_embed=true reuse_example_batch=1\ + use_iota_embed=true reuse_example_batch=1\ dataset_type=synthetic attention='flash' gcs_metrics=true diff --git a/src/maxtext/configs/tpu/v5e/32b.sh b/src/maxtext/configs/tpu/v5e/32b.sh index 2e72b5da2e..ff870948c3 100644 --- a/src/maxtext/configs/tpu/v5e/32b.sh +++ b/src/maxtext/configs/tpu/v5e/32b.sh @@ -4,11 +4,10 @@ echo "Running 32b.sh" # # Command Flags: # OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml) -# DATASET_PATH (Required, unless dataset_path is already set in base.yml) # RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE) # # Example to invoke this script: -# bash src/maxtext/configs/tpu/v5e/32b.sh RUN_NAME="" OUTPUT_PATH="gs://" DATASET_PATH="gs://" +# bash src/maxtext/configs/tpu/v5e/32b.sh RUN_NAME="" OUTPUT_PATH="gs://" # # Example to AOT compile: # bash src/maxtext/configs/tpu/v5e/32b.sh EXECUTABLE=train_compile M_COMPILE_TOPOLOGY=v5e-256 M_COMPILE_TOPOLOGY_NUM_SLICES=2 @@ -45,5 +44,5 @@ python3 -m maxtext.trainers.pre_train.$EXECUTABLE "${MAXTEXT_CONFIGS_DIR:-${MAXT steps=15 per_device_batch_size=4 enable_checkpointing=false\ remat_policy=full global_parameter_scale=32\ max_target_length=2048 base_output_directory=$OUTPUT_PATH\ - dataset_path=$DATASET_PATH use_iota_embed=true reuse_example_batch=1\ + use_iota_embed=true reuse_example_batch=1\ dataset_type=synthetic attention='flash' gcs_metrics=true diff --git a/src/maxtext/configs/tpu/v5e/64b.sh b/src/maxtext/configs/tpu/v5e/64b.sh index f28eeb6394..faec50b553 100644 --- a/src/maxtext/configs/tpu/v5e/64b.sh +++ b/src/maxtext/configs/tpu/v5e/64b.sh @@ -4,11 +4,10 @@ echo "Running 64b.sh" # # Command Flags: # OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml) -# DATASET_PATH (Required, unless dataset_path is already set in base.yml) # RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE) # # Example to invoke this script: -# bash src/maxtext/configs/tpu/v5e/64b.sh RUN_NAME="" OUTPUT_PATH="gs://" DATASET_PATH="gs://" +# bash src/maxtext/configs/tpu/v5e/64b.sh RUN_NAME="" OUTPUT_PATH="gs://" # # Example to AOT compile: # bash src/maxtext/configs/tpu/v5e/64b.sh EXECUTABLE=train_compile M_COMPILE_TOPOLOGY=v5e-256 M_COMPILE_TOPOLOGY_NUM_SLICES=2 @@ -45,5 +44,5 @@ python3 -m maxtext.trainers.pre_train.$EXECUTABLE "${MAXTEXT_CONFIGS_DIR:-${MAXT steps=15 per_device_batch_size=2 enable_checkpointing=false\ remat_policy=full global_parameter_scale=64\ max_target_length=2048 base_output_directory=$OUTPUT_PATH\ - dataset_path=$DATASET_PATH use_iota_embed=true reuse_example_batch=1\ + use_iota_embed=true reuse_example_batch=1\ dataset_type=synthetic attention='flash' gcs_metrics=true diff --git a/src/maxtext/configs/tpu/v5e/gpt3_175b.sh b/src/maxtext/configs/tpu/v5e/gpt3_175b.sh index d71f9b1a78..d9c95ce2c6 100644 --- a/src/maxtext/configs/tpu/v5e/gpt3_175b.sh +++ b/src/maxtext/configs/tpu/v5e/gpt3_175b.sh @@ -3,11 +3,10 @@ # # Command Flags: # OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml) -# DATASET_PATH (Required, unless dataset_path is already set in base.yml) # RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE) # # Example to invoke this script: -# bash src/maxtext/configs/tpu/v5e/gpt3_175b.sh RUN_NAME="" OUTPUT_PATH="gs://" DATASET_PATH="gs://" +# bash src/maxtext/configs/tpu/v5e/gpt3_175b.sh RUN_NAME="" OUTPUT_PATH="gs://" # # Example to AOT compile: # bash src/maxtext/configs/tpu/v5e/gpt3_175b.sh EXECUTABLE=train_compile M_COMPILE_TOPOLOGY=v5e-256 M_COMPILE_TOPOLOGY_NUM_SLICES=8 diff --git a/src/maxtext/configs/tpu/v5e/llama2_13b.sh b/src/maxtext/configs/tpu/v5e/llama2_13b.sh index 8dca05f36f..e79997da64 100644 --- a/src/maxtext/configs/tpu/v5e/llama2_13b.sh +++ b/src/maxtext/configs/tpu/v5e/llama2_13b.sh @@ -3,11 +3,10 @@ # # Command Flags: # OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml) -# DATASET_PATH (Required, unless dataset_path is already set in base.yml) # RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE) # # Example to invoke this script: -# bash src/maxtext/configs/tpu/v5e/llama2_13b.sh RUN_NAME="" OUTPUT_PATH="gs://" DATASET_PATH="gs://" +# bash src/maxtext/configs/tpu/v5e/llama2_13b.sh RUN_NAME="" OUTPUT_PATH="gs://" # # Example to AOT compile: # bash src/maxtext/configs/tpu/v5e/llama2_13b.sh EXECUTABLE=train_compile M_COMPILE_TOPOLOGY=v5e-256 M_COMPILE_TOPOLOGY_NUM_SLICES=2 @@ -42,6 +41,6 @@ fi export LIBTPU_INIT_ARGS="--xla_tpu_enable_data_parallel_all_reduce_opt=true --xla_tpu_data_parallel_opt_different_sized_ops=true --xla_tpu_enable_async_collective_fusion=true --xla_tpu_enable_async_collective_fusion_fuse_all_gather=true --xla_tpu_enable_async_collective_fusion_multiple_steps=true --xla_tpu_overlap_compute_collective_tc=true --xla_enable_async_all_gather=true" python3 -m maxtext.trainers.pre_train.$EXECUTABLE "${MAXTEXT_CONFIGS_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/maxtext/configs}"//base.yml model_name=llama2-13b\ - base_output_directory=$OUTPUT_PATH dataset_path=${DATASET_PATH}\ + base_output_directory=$OUTPUT_PATH dataset_type=synthetic\ tokenizer_path="${MAXTEXT_ASSETS_ROOT:-${MAXTEXT_PKG_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/maxtext/assets/tokenizers}}"/tokenizer.llama2 per_device_batch_size=8 remat_policy=qkv_proj_offloaded\ steps=15 enable_checkpointing=false use_iota_embed=true diff --git a/src/maxtext/configs/tpu/v5e/llama2_70b.sh b/src/maxtext/configs/tpu/v5e/llama2_70b.sh index 528970ccc2..aee3babe44 100644 --- a/src/maxtext/configs/tpu/v5e/llama2_70b.sh +++ b/src/maxtext/configs/tpu/v5e/llama2_70b.sh @@ -3,11 +3,10 @@ # # Command Flags: # OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml) -# DATASET_PATH (Required, unless dataset_path is already set in base.yml) # RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE) # # Example to invoke this script: -# bash src/maxtext/configs/tpu/v5e/llama2_70b.sh RUN_NAME="" OUTPUT_PATH="gs://" DATASET_PATH="gs://" +# bash src/maxtext/configs/tpu/v5e/llama2_70b.sh RUN_NAME="" OUTPUT_PATH="gs://" # # Example to AOT compile: # bash src/maxtext/configs/tpu/v5e/llama2_70b.sh EXECUTABLE=train_compile M_COMPILE_TOPOLOGY=v5e-256 M_COMPILE_TOPOLOGY_NUM_SLICES=2 @@ -42,6 +41,6 @@ fi export LIBTPU_INIT_ARGS="--xla_tpu_enable_data_parallel_all_reduce_opt=true --xla_tpu_data_parallel_opt_different_sized_ops=true --xla_tpu_enable_async_collective_fusion=true --xla_tpu_enable_async_collective_fusion_fuse_all_gather=true --xla_tpu_enable_async_collective_fusion_multiple_steps=true --xla_tpu_overlap_compute_collective_tc=true --xla_enable_async_all_gather=true" python3 -m maxtext.trainers.pre_train.$EXECUTABLE "${MAXTEXT_CONFIGS_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/maxtext/configs}"//base.yml model_name=llama2-70b\ - base_output_directory=$OUTPUT_PATH dataset_path=${DATASET_PATH}\ + base_output_directory=$OUTPUT_PATH dataset_type=synthetic\ tokenizer_path="${MAXTEXT_ASSETS_ROOT:-${MAXTEXT_PKG_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/maxtext/assets/tokenizers}}"/tokenizer.llama2 per_device_batch_size=2 remat_policy=qkv_proj_offloaded\ steps=15 enable_checkpointing=false use_iota_embed=true diff --git a/src/maxtext/configs/tpu/v5e/llama2_7b.sh b/src/maxtext/configs/tpu/v5e/llama2_7b.sh index 6b68297a1e..e4aa662b2c 100644 --- a/src/maxtext/configs/tpu/v5e/llama2_7b.sh +++ b/src/maxtext/configs/tpu/v5e/llama2_7b.sh @@ -3,11 +3,10 @@ # # Command Flags: # OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml) -# DATASET_PATH (Required, unless dataset_path is already set in base.yml) # RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE) # # Example to invoke this script: -# bash src/maxtext/configs/tpu/v5e/llama2_7b.sh RUN_NAME="" OUTPUT_PATH="gs://" DATASET_PATH="gs://" +# bash src/maxtext/configs/tpu/v5e/llama2_7b.sh RUN_NAME="" OUTPUT_PATH="gs://" # # Example to AOT compile: # bash src/maxtext/configs/tpu/v5e/llama2_7b.sh EXECUTABLE=train_compile M_COMPILE_TOPOLOGY=v5e-256 M_COMPILE_TOPOLOGY_NUM_SLICES=2 @@ -42,6 +41,6 @@ fi export LIBTPU_INIT_ARGS="--xla_tpu_enable_data_parallel_all_reduce_opt=true --xla_tpu_data_parallel_opt_different_sized_ops=true --xla_tpu_enable_async_collective_fusion=true --xla_tpu_enable_async_collective_fusion_fuse_all_gather=true --xla_tpu_enable_async_collective_fusion_multiple_steps=true --xla_tpu_overlap_compute_collective_tc=true --xla_enable_async_all_gather=true" python3 -m maxtext.trainers.pre_train.$EXECUTABLE "${MAXTEXT_CONFIGS_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/maxtext/configs}"//base.yml model_name=llama2-7b\ - base_output_directory=$OUTPUT_PATH dataset_path=${DATASET_PATH}\ + base_output_directory=$OUTPUT_PATH dataset_type=synthetic\ tokenizer_path="${MAXTEXT_ASSETS_ROOT:-${MAXTEXT_PKG_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/maxtext/assets/tokenizers}}"/tokenizer.llama2 per_device_batch_size=4 remat_policy=save_qkv_proj\ steps=15 enable_checkpointing=false use_iota_embed=true \ No newline at end of file diff --git a/src/maxtext/configs/tpu/v5p/1024b.sh b/src/maxtext/configs/tpu/v5p/1024b.sh index 14bdcc6692..d42bc8a70e 100644 --- a/src/maxtext/configs/tpu/v5p/1024b.sh +++ b/src/maxtext/configs/tpu/v5p/1024b.sh @@ -4,11 +4,10 @@ echo "Running 1024b.sh" # # Command Flags: # OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml) -# DATASET_PATH (Required, unless dataset_path is already set in base.yml) # RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE) # # Example to invoke this script: -# bash src/maxtext/configs/tpu/v5p/1024b.sh RUN_NAME="" OUTPUT_PATH="gs://" DATASET_PATH="gs://" +# bash src/maxtext/configs/tpu/v5p/1024b.sh RUN_NAME="" OUTPUT_PATH="gs://" # # Example to AOT compile: # bash src/maxtext/configs/tpu/v5p/1024b.sh EXECUTABLE=train_compile M_COMPILE_TOPOLOGY=v5p-2048 M_COMPILE_TOPOLOGY_NUM_SLICES=2 @@ -46,5 +45,5 @@ python3 -m maxtext.trainers.pre_train.$EXECUTABLE "${MAXTEXT_CONFIGS_DIR:-${MAXT remat_policy=full global_parameter_scale=1024\ ici_fsdp_parallelism=-1 ici_tensor_parallelism=16\ max_target_length=2048 base_output_directory=$OUTPUT_PATH\ - dataset_path=$DATASET_PATH use_iota_embed=true reuse_example_batch=1\ + use_iota_embed=true reuse_example_batch=1\ dataset_type=synthetic gcs_metrics=true attention='flash' quantization="" diff --git a/src/maxtext/configs/tpu/v5p/128b.sh b/src/maxtext/configs/tpu/v5p/128b.sh index ce068f7fde..e53c6a9ede 100644 --- a/src/maxtext/configs/tpu/v5p/128b.sh +++ b/src/maxtext/configs/tpu/v5p/128b.sh @@ -4,11 +4,10 @@ echo "Running 128b.sh" # # Command Flags: # OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml) -# DATASET_PATH (Required, unless dataset_path is already set in base.yml) # RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE) # # Example to invoke this script: -# bash src/maxtext/configs/tpu/v5p/128b.sh RUN_NAME="" OUTPUT_PATH="gs://" DATASET_PATH="gs://" +# bash src/maxtext/configs/tpu/v5p/128b.sh RUN_NAME="" OUTPUT_PATH="gs://" # # Example to AOT compile: # bash src/maxtext/configs/tpu/v5p/128b.sh EXECUTABLE=train_compile M_COMPILE_TOPOLOGY=v5p-256 M_COMPILE_TOPOLOGY_NUM_SLICES=2 @@ -46,5 +45,5 @@ python3 -m maxtext.trainers.pre_train.$EXECUTABLE "${MAXTEXT_CONFIGS_DIR:-${MAXT remat_policy=minimal global_parameter_scale=128\ ici_fsdp_parallelism=-1 ici_tensor_parallelism=8\ max_target_length=2048 base_output_directory=$OUTPUT_PATH\ - dataset_path=$DATASET_PATH use_iota_embed=true reuse_example_batch=1\ + use_iota_embed=true reuse_example_batch=1\ dataset_type=synthetic gcs_metrics=true attention='flash' quantization="" diff --git a/src/maxtext/configs/tpu/v5p/256b.sh b/src/maxtext/configs/tpu/v5p/256b.sh index 04e16ca0cc..984f281e32 100644 --- a/src/maxtext/configs/tpu/v5p/256b.sh +++ b/src/maxtext/configs/tpu/v5p/256b.sh @@ -4,12 +4,11 @@ echo "Running 256b.sh" # # Command Flags: # OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml) -# DATASET_PATH (Required, unless dataset_path is already set in base.yml) # RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE) # PLATFORM (Optional, can be "gke" or "gce", default is "gce") # # Example to invoke this script: -# bash src/maxtext/configs/tpu/v5p/256b.sh RUN_NAME="" OUTPUT_PATH="gs://" DATASET_PATH="gs://" +# bash src/maxtext/configs/tpu/v5p/256b.sh RUN_NAME="" OUTPUT_PATH="gs://" # # Example to AOT compile: # bash src/maxtext/configs/tpu/v5p/256b.sh EXECUTABLE=train_compile M_COMPILE_TOPOLOGY=v5p-1024 M_COMPILE_TOPOLOGY_NUM_SLICES=2 @@ -47,5 +46,5 @@ python3 -m maxtext.trainers.pre_train.$EXECUTABLE "${MAXTEXT_CONFIGS_DIR:-${MAXT remat_policy=minimal global_parameter_scale=256\ ici_fsdp_parallelism=-1 ici_tensor_parallelism=8\ max_target_length=2048 base_output_directory=$OUTPUT_PATH\ - dataset_path=$DATASET_PATH use_iota_embed=true reuse_example_batch=1\ + use_iota_embed=true reuse_example_batch=1\ dataset_type=synthetic gcs_metrics=true attention='flash' quantization="" diff --git a/src/maxtext/configs/tpu/v5p/32b.sh b/src/maxtext/configs/tpu/v5p/32b.sh index 813f05d3a8..de129f0ec5 100644 --- a/src/maxtext/configs/tpu/v5p/32b.sh +++ b/src/maxtext/configs/tpu/v5p/32b.sh @@ -4,11 +4,10 @@ echo "Running 32b.sh" # # Command Flags: # OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml) -# DATASET_PATH (Required, unless dataset_path is already set in base.yml) # RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE) # # Example to invoke this script: -# bash src/maxtext/configs/tpu/v5p/32b.sh RUN_NAME="" OUTPUT_PATH="gs://" DATASET_PATH="gs://" +# bash src/maxtext/configs/tpu/v5p/32b.sh RUN_NAME="" OUTPUT_PATH="gs://" # # Example to AOT compile: # bash src/maxtext/configs/tpu/v5p/32b.sh EXECUTABLE=train_compile M_COMPILE_TOPOLOGY=v5p-128 M_COMPILE_TOPOLOGY_NUM_SLICES=2 @@ -46,5 +45,5 @@ python3 -m maxtext.trainers.pre_train.$EXECUTABLE "${MAXTEXT_CONFIGS_DIR:-${MAXT remat_policy=minimal global_parameter_scale=32\ ici_fsdp_parallelism=-1 ici_tensor_parallelism=4\ max_target_length=2048 base_output_directory=$OUTPUT_PATH\ - dataset_path=$DATASET_PATH use_iota_embed=true reuse_example_batch=1\ + use_iota_embed=true reuse_example_batch=1\ dataset_type=synthetic gcs_metrics=true attention='flash' quantization="" diff --git a/src/maxtext/configs/tpu/v5p/512b.sh b/src/maxtext/configs/tpu/v5p/512b.sh index 293c0b65eb..42e6c7cec7 100644 --- a/src/maxtext/configs/tpu/v5p/512b.sh +++ b/src/maxtext/configs/tpu/v5p/512b.sh @@ -4,12 +4,11 @@ echo "Running 512b.sh" # # Command Flags: # OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml) -# DATASET_PATH (Required, unless dataset_path is already set in base.yml) # RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE) # PLATFORM (Optional, can be "gke" or "gce", default is "gce") # # Example to invoke this script: -# bash src/maxtext/configs/tpu/v5p/512b.sh RUN_NAME="" OUTPUT_PATH="gs://" DATASET_PATH="gs://" +# bash src/maxtext/configs/tpu/v5p/512b.sh RUN_NAME="" OUTPUT_PATH="gs://" # # Example to AOT compile: # bash src/maxtext/configs/tpu/v5p/512b.sh EXECUTABLE=train_compile M_COMPILE_TOPOLOGY=v5p-1024 M_COMPILE_TOPOLOGY_NUM_SLICES=2 @@ -47,5 +46,5 @@ python3 -m maxtext.trainers.pre_train.$EXECUTABLE "${MAXTEXT_CONFIGS_DIR:-${MAXT remat_policy=full global_parameter_scale=512\ ici_fsdp_parallelism=-1 ici_tensor_parallelism=8\ max_target_length=2048 base_output_directory=$OUTPUT_PATH\ - dataset_path=$DATASET_PATH use_iota_embed=true reuse_example_batch=1\ + use_iota_embed=true reuse_example_batch=1\ dataset_type=synthetic gcs_metrics=true attention='flash' quantization="" diff --git a/src/maxtext/configs/tpu/v5p/64b.sh b/src/maxtext/configs/tpu/v5p/64b.sh index ec44e130f3..90c01c2c74 100644 --- a/src/maxtext/configs/tpu/v5p/64b.sh +++ b/src/maxtext/configs/tpu/v5p/64b.sh @@ -4,11 +4,10 @@ echo "Running 64b.sh" # # Command Flags: # OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml) -# DATASET_PATH (Required, unless dataset_path is already set in base.yml) # RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE) # # Example to invoke this script: -# bash src/maxtext/configs/tpu/v5p/64b.sh RUN_NAME="" OUTPUT_PATH="gs://" DATASET_PATH="gs://" +# bash src/maxtext/configs/tpu/v5p/64b.sh RUN_NAME="" OUTPUT_PATH="gs://" # # Example to AOT compile: # bash src/maxtext/configs/tpu/v5p/64b.sh EXECUTABLE=train_compile M_COMPILE_TOPOLOGY=v5p-128 M_COMPILE_TOPOLOGY_NUM_SLICES=2 @@ -46,5 +45,5 @@ python3 -m maxtext.trainers.pre_train.$EXECUTABLE "${MAXTEXT_CONFIGS_DIR:-${MAXT remat_policy=minimal global_parameter_scale=64\ ici_fsdp_parallelism=-1 ici_tensor_parallelism=4\ max_target_length=2048 base_output_directory=$OUTPUT_PATH\ - dataset_path=$DATASET_PATH use_iota_embed=true reuse_example_batch=1\ + use_iota_embed=true reuse_example_batch=1\ dataset_type=synthetic gcs_metrics=true attention='flash' quantization="" diff --git a/src/maxtext/configs/tpu/v5p/llama2_70b.sh b/src/maxtext/configs/tpu/v5p/llama2_70b.sh index b0e1836ea1..2228794617 100644 --- a/src/maxtext/configs/tpu/v5p/llama2_70b.sh +++ b/src/maxtext/configs/tpu/v5p/llama2_70b.sh @@ -5,7 +5,6 @@ # # Command Flags: # OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml) -# DATASET_PATH (Required, unless dataset_path is already set in base.yml) # RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE) # # Example to invoke this script: @@ -45,7 +44,7 @@ fi export LIBTPU_INIT_ARGS="--xla_tpu_enable_async_collective_fusion_fuse_all_gather=true --xla_tpu_megacore_fusion_allow_ags=false --xla_enable_async_collective_permute=true --xla_tpu_enable_ag_backward_pipelining=true --xla_tpu_enable_data_parallel_all_reduce_opt=true --xla_tpu_data_parallel_opt_different_sized_ops=true --xla_tpu_enable_async_collective_fusion=true --xla_tpu_enable_async_collective_fusion_multiple_steps=true --xla_tpu_overlap_compute_collective_tc=true --xla_enable_async_all_gather=true" python3 -m maxtext.trainers.pre_train.$EXECUTABLE "${MAXTEXT_CONFIGS_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/maxtext/configs}"//base.yml model_name=llama2-70b\ - base_output_directory=$OUTPUT_PATH dataset_path=${DATASET_PATH}\ + base_output_directory=$OUTPUT_PATH\ tokenizer_path="${MAXTEXT_ASSETS_ROOT:-${MAXTEXT_PKG_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/maxtext/assets/tokenizers}}"/tokenizer.llama2 remat_policy=save_dot_except_mlpwi per_device_batch_size=4\ steps=30 enable_checkpointing=false use_iota_embed=true max_target_length=4096\ profiler=xplane skip_first_n_steps_for_profiler=10 profiler_steps=5 gcs_metrics=true\ diff --git a/src/maxtext/configs/tpu/v5p/llama2_7b.sh b/src/maxtext/configs/tpu/v5p/llama2_7b.sh index 08683ace65..774b7ffe61 100644 --- a/src/maxtext/configs/tpu/v5p/llama2_7b.sh +++ b/src/maxtext/configs/tpu/v5p/llama2_7b.sh @@ -5,7 +5,6 @@ # # Command Flags: # OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml) -# DATASET_PATH (Required, unless dataset_path is already set in base.yml or using synthetic dataset_type) # RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE) # # Example to invoke this script: @@ -45,7 +44,7 @@ fi # Train export LIBTPU_INIT_ARGS="--xla_tpu_enable_async_collective_fusion_fuse_all_gather=true --xla_tpu_megacore_fusion_allow_ags=false --xla_enable_async_collective_permute=true --xla_tpu_enable_ag_backward_pipelining=true --xla_tpu_enable_data_parallel_all_reduce_opt=true --xla_tpu_data_parallel_opt_different_sized_ops=true --xla_tpu_enable_async_collective_fusion=true --xla_tpu_enable_async_collective_fusion_multiple_steps=true --xla_tpu_overlap_compute_collective_tc=true --xla_enable_async_all_gather=true" python3 -m maxtext.trainers.pre_train.$EXECUTABLE "${MAXTEXT_CONFIGS_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/maxtext/configs}"//base.yml model_name=llama2-7b\ - base_output_directory=$OUTPUT_PATH dataset_path=${DATASET_PATH}\ + base_output_directory=$OUTPUT_PATH\ tokenizer_path="${MAXTEXT_ASSETS_ROOT:-${MAXTEXT_PKG_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/maxtext/assets/tokenizers}}"/tokenizer.llama2 remat_policy=minimal per_device_batch_size=4\ steps=30 enable_checkpointing=false use_iota_embed=true max_target_length=4096\ profiler=xplane skip_first_n_steps_for_profiler=10 profiler_steps=5 gcs_metrics=true\ diff --git a/src/maxtext/configs/tpu/v6e/gemma2_27b.sh b/src/maxtext/configs/tpu/v6e/gemma2_27b.sh index 4bbc58643e..970722c6f8 100644 --- a/src/maxtext/configs/tpu/v6e/gemma2_27b.sh +++ b/src/maxtext/configs/tpu/v6e/gemma2_27b.sh @@ -3,11 +3,10 @@ # # Command Flags: # OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml) -# DATASET_PATH (Required, unless dataset_path is already set in base.yml) # RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE) # # Example to invoke this script: -# bash src/maxtext/configs/tpu/v6e/gemma2_27b.sh RUN_NAME="" OUTPUT_PATH="gs://" DATASET_PATH="gs://" +# bash src/maxtext/configs/tpu/v6e/gemma2_27b.sh RUN_NAME="" OUTPUT_PATH="gs://" # diff --git a/src/maxtext/configs/tpu/v6e/gemma2_9b.sh b/src/maxtext/configs/tpu/v6e/gemma2_9b.sh index 2a7deb01aa..2728f46aed 100644 --- a/src/maxtext/configs/tpu/v6e/gemma2_9b.sh +++ b/src/maxtext/configs/tpu/v6e/gemma2_9b.sh @@ -3,11 +3,10 @@ # # Command Flags: # OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml) -# DATASET_PATH (Required, unless dataset_path is already set in base.yml) # RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE) # # Example to invoke this script: -# bash src/maxtext/configs/tpu/v6e/gemma2_9b.sh RUN_NAME="" OUTPUT_PATH="gs://" DATASET_PATH="gs://" +# bash src/maxtext/configs/tpu/v6e/gemma2_9b.sh RUN_NAME="" OUTPUT_PATH="gs://" # diff --git a/src/maxtext/configs/tpu/v6e/gemma3_27b.sh b/src/maxtext/configs/tpu/v6e/gemma3_27b.sh index 7441dc358a..80824e9f46 100644 --- a/src/maxtext/configs/tpu/v6e/gemma3_27b.sh +++ b/src/maxtext/configs/tpu/v6e/gemma3_27b.sh @@ -3,11 +3,10 @@ # # Command Flags: # OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml) -# DATASET_PATH (Required, unless dataset_path is already set in base.yml) # RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE) # # Example to invoke this script: -# bash src/maxtext/configs/tpu/v6e/gemma3_27b.sh RUN_NAME="" OUTPUT_PATH="gs://" DATASET_PATH="gs://" +# bash src/maxtext/configs/tpu/v6e/gemma3_27b.sh RUN_NAME="" OUTPUT_PATH="gs://" # diff --git a/src/maxtext/configs/tpu/v6e/gpt3_175b.sh b/src/maxtext/configs/tpu/v6e/gpt3_175b.sh index 79ca9d30fc..b7f3b40454 100644 --- a/src/maxtext/configs/tpu/v6e/gpt3_175b.sh +++ b/src/maxtext/configs/tpu/v6e/gpt3_175b.sh @@ -3,11 +3,10 @@ # # Command Flags: # OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml) -# DATASET_PATH (Required, unless dataset_path is already set in base.yml) # RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE) # # Example to invoke this script: -# bash src/maxtext/configs/tpu/v6e/gpt3_175b.sh RUN_NAME="" OUTPUT_PATH="gs://" DATASET_PATH="gs://" +# bash src/maxtext/configs/tpu/v6e/gpt3_175b.sh RUN_NAME="" OUTPUT_PATH="gs://" # diff --git a/src/maxtext/configs/tpu/v6e/llama2_7b_4096.sh b/src/maxtext/configs/tpu/v6e/llama2_7b_4096.sh index 4e08e3d76f..8279acede6 100644 --- a/src/maxtext/configs/tpu/v6e/llama2_7b_4096.sh +++ b/src/maxtext/configs/tpu/v6e/llama2_7b_4096.sh @@ -3,11 +3,10 @@ # # Command Flags: # OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml) -# DATASET_PATH (Required, unless dataset_path is already set in base.yml) # RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE) # # Example to invoke this script: -# bash src/maxtext/configs/tpu/v6e/llama2_7b_4096.sh RUN_NAME="" OUTPUT_PATH="gs://" DATASET_PATH="gs://" +# bash src/maxtext/configs/tpu/v6e/llama2_7b_4096.sh RUN_NAME="" OUTPUT_PATH="gs://" # diff --git a/src/maxtext/configs/tpu/v6e/mixtral_8x7b.sh b/src/maxtext/configs/tpu/v6e/mixtral_8x7b.sh index 2eac06d68c..c5bd75d94f 100644 --- a/src/maxtext/configs/tpu/v6e/mixtral_8x7b.sh +++ b/src/maxtext/configs/tpu/v6e/mixtral_8x7b.sh @@ -3,11 +3,10 @@ # # Command Flags: # OUTPUT_PATH (Required, unless base_output_directory is already set in base.yml) -# DATASET_PATH (Required, unless dataset_path is already set in base.yml) # RUN_NAME (Required, unless run_name is already set in base.yml or running with XPK/GKE) # # Example to invoke this script: -# bash src/maxtext/configs/tpu/v6e/mixtral_8x7b.sh RUN_NAME="" OUTPUT_PATH="gs://" DATASET_PATH="gs://" +# bash src/maxtext/configs/tpu/v6e/mixtral_8x7b.sh RUN_NAME="" OUTPUT_PATH="gs://" # diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index b0f407eff3..3dffe5e7ac 100644 --- a/src/maxtext/configs/types.py +++ b/src/maxtext/configs/types.py @@ -470,6 +470,9 @@ class Quantization(BaseModel): 50, description=("The first number of steps before updating the sparsity masks."), ) + use_te_comm_gemm_overlap: bool = Field( + False, description="If True, uses Transformer Engine's collective GEMM overlap algorithm." + ) class ModelArchitecture(BaseModel): @@ -963,7 +966,6 @@ class HardwareAndMesh(BaseModel): "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive", @@ -972,7 +974,13 @@ class HardwareAndMesh(BaseModel): ) shard_mode: ShardMode = Field("auto", description="can be either auto or explicit") inhomogeneous_layer_cycle_interval: int = Field(1, description="The interval of repeated inhomogeneous layer patterns.") - scan_layers: bool = Field(True, description="Whether to use jax.lax.scan over layers.") + scan_layers: bool = Field( + True, + description=( + "Whether to use jax.lax.scan over layers (stacked/unstacked checkpoint). " + "When resuming from a checkpoint, this flag is auto-determined from metadata." + ), + ) param_scan_axis: int = Field(1, description="Axis to scan over for parameters.") context_parallel_load_balance: bool = Field(True, description="Whether to use load balancing for context parallelism.") context_parallel_strategy: str = Field( @@ -988,11 +996,11 @@ class HardwareAndMesh(BaseModel): CustomRule.DEFAULT, description="Customized mesh and logical rules for granularity." ) allow_split_physical_axes: bool = Field(False, description="Allow splitting physical axes for device mesh creation.") - enable_nnx: bool = Field(False, description="Whether to use NNX for model definition.") + enable_nnx: bool = Field(True, description="Whether to use NNX for model definition.") optimize_mesh_for_tpu_v6e: bool = Field(False, description="Apply transformations to the mesh for TPU v6e.") shardy: bool = Field(True, description="Whether to use shardy XLA backend.") - pure_nnx_decoder: bool = Field(False, description="Whether to enable pure NNX decoder.") - pure_nnx: bool = Field(False, description="Whether to enable pure NNX mode.") + pure_nnx_decoder: bool = Field(True, description="Whether to enable pure NNX decoder.") + pure_nnx: bool = Field(True, description="Whether to enable pure NNX mode.") remove_size_one_mesh_axis_from_type: bool = Field( True, description="Whether to remove size one mesh axis from type through jax.config." ) @@ -1036,7 +1044,6 @@ class DcnParallelism(BaseModel): dcn_context_parallelism: int = Field(1, description="DCN axis for context parallelism.") dcn_context_autoregressive_parallelism: int = Field(1, description="DCN axis for context autoregressive parallelism.") dcn_tensor_parallelism: int = Field(1, description="DCN axis for tensor parallelism (not recommended).") - dcn_tensor_transpose_parallelism: int = Field(1, description="DCN axis for tensor transpose parallelism.") dcn_tensor_sequence_parallelism: int = Field( 1, description="DCN axis for tensor sequence parallelism (not recommended)." ) @@ -1056,7 +1063,6 @@ class IciParallelism(BaseModel): ici_context_parallelism: int = Field(1, description="ICI axis for context parallelism.") ici_context_autoregressive_parallelism: int = Field(1, description="ICI axis for context autoregressive parallelism.") ici_tensor_parallelism: int = Field(1, description="ICI axis for tensor parallelism.") - ici_tensor_transpose_parallelism: int = Field(1, description="ICI axis for tensor transpose parallelism.") ici_tensor_sequence_parallelism: int = Field(1, description="ICI axis for tensor sequence parallelism.") ici_autoregressive_parallelism: int = Field(1, description="ICI axis for autoregressive parallelism.") ici_pipeline_parallelism: int = Field(1, description="ICI axis for pipeline parallelism.") @@ -2156,10 +2162,6 @@ class RL(BaseModel): "", description="System prompt injected into the agent at rollout time (agentic only).", ) - degenerate_group_masking: bool = Field( - True, - description="Mask degenerate groups (all-zero advantages) from contributing to loss (agentic only).", - ) epsilon_high: Optional[float] = Field( None, description="Upper-bound clipping epsilon for GRPO loss. Defaults to epsilon when None (agentic only).", @@ -2191,6 +2193,20 @@ class RLDataset(BaseModel): "When set, replaces the built-in dataset processor for custom datasets." ), ) + reward_functions_path: str = Field( + "", + description=( + "Optional path to a user Python file containing custom reward functions. " + "Used with `reward_functions` to fully replace the built-in reward stack." + ), + ) + reward_functions: str = Field( + "", + description=( + "Comma-separated names of reward functions to import from `reward_functions_path`. " + "Each function signature: (prompts, completions, tmvp_config, **kwargs) -> list[float]." + ), + ) class RLEvaluation(BaseModel): @@ -2555,6 +2571,26 @@ def validate_ragged_buffer_factor(self): " 2. Ragged sort with ring of experts (use_ring_of_experts=True AND use_ragged_sort=True)" ) + def _validate_use_te_comm_gemm_overlap(self): + """Validates that use_te_comm_gemm_overlap is used with supported settings to enable TE Collective GEMM ops.""" + te_has_distributed_env = jax.local_device_count() == 1 and jax.distributed.is_initialized() + + if self.hardware != "gpu_multiprocess" or not te_has_distributed_env: + raise ValueError( + "TE Collective GEMM operations are only supported for hardware=gpu_multiprocess with rank per GPU." + ) + + tsp_size = self.ici_tensor_sequence_parallelism * self.dcn_tensor_sequence_parallelism + tp_size = self.ici_tensor_parallelism * self.dcn_tensor_parallelism + + if tsp_size == 1 or tp_size > 1: + raise ValueError("TE Collective GEMM operations are only supported for TSP size > 1 and TP size == 1.") + + if not self.quantization.startswith("te_"): + raise ValueError( + "TE Collective GEMM operations are only supported for TE quantization recipes (i.e. starting with 'te_')." + ) + @model_validator(mode="after") def set_derived_and_validate_values(self) -> "MaxTextConfig": """ @@ -3393,7 +3429,6 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de "context": self.ici_context_parallelism, "context_autoregressive": self.ici_context_autoregressive_parallelism, "tensor": self.ici_tensor_parallelism, - "tensor_transpose": self.ici_tensor_transpose_parallelism, "tensor_sequence": self.ici_tensor_sequence_parallelism, "model": self.ici_tensor_parallelism, "expert": self.ici_expert_parallelism, @@ -3413,7 +3448,6 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de "context": self.dcn_context_parallelism, "context_autoregressive": self.dcn_context_autoregressive_parallelism, "tensor": self.dcn_tensor_parallelism, - "tensor_transpose": self.dcn_tensor_transpose_parallelism, "tensor_sequence": self.dcn_tensor_sequence_parallelism, "model": self.dcn_tensor_parallelism, "expert": self.dcn_expert_parallelism, @@ -3453,6 +3487,9 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de self._validate_check_vma_is_supported() + if self.use_te_comm_gemm_overlap: + self._validate_use_te_comm_gemm_overlap() + # Final string-to-enum conversions if they haven't been coerced by pydantic yet. if isinstance(self.decoder_block, str): self.decoder_block = DecoderBlockType(self.decoder_block.lower()) diff --git a/src/maxtext/examples/chat_templates/gpt_oss_rl.json b/src/maxtext/examples/chat_templates/gpt_oss_rl.json new file mode 100644 index 0000000000..9f6d522c43 --- /dev/null +++ b/src/maxtext/examples/chat_templates/gpt_oss_rl.json @@ -0,0 +1,4 @@ +{ + "SYSTEM_PROMPT": "You are given a problem. Think about the problem and provide your reasoning. Place it between {reasoning_start_token} and {reasoning_end_token}. Then, provide the final answer (i.e., just one numerical value) between {solution_start_token} and {solution_end_token}.", + "TEMPLATE": "<|start|>system<|message|>You are ChatGPT, a large language model trained by OpenAI.\nKnowledge cutoff: 2024-06\nCurrent date: 2026-06-19\n\nReasoning: medium\n\n# Valid channels: analysis, commentary, final. Channel must be included for every message.\nCalls to these tools must go to the commentary channel: 'functions'.<|end|><|start|>user<|message|>user\n{system_prompt}\n\n{question}\nmodel<|end|><|start|>assistant" +} diff --git a/src/maxtext/experimental/rl/grpo_trainer.py b/src/maxtext/experimental/rl/grpo_trainer.py index 9ea0a68f82..e1d80e8917 100644 --- a/src/maxtext/experimental/rl/grpo_trainer.py +++ b/src/maxtext/experimental/rl/grpo_trainer.py @@ -864,7 +864,7 @@ def init_state_fn(): with maybe_record_goodput(recorder, GoodputEvent.TRAINING_PREPARATION): data_iterator = grpo_input_pipeline.create_data_iterator(config_inference, inference_mesh) - state, _, state_mesh_shardings, data_iterator = maxtext_utils.setup_training_state( + state, _, state_mesh_shardings, data_iterator, _ = maxtext_utils.setup_training_state( data_iterator, config, mesh, checkpoint_manager, init_state_fn ) diff --git a/src/maxtext/inference/decode.py b/src/maxtext/inference/decode.py index 68cd450c99..9b7f415c70 100644 --- a/src/maxtext/inference/decode.py +++ b/src/maxtext/inference/decode.py @@ -139,7 +139,7 @@ def main(argv: Sequence[str]) -> None: if config.use_multimodal: tokens = mm_processor.prepare_text_for_image_fusion(tokens=tokens, config=config, processor_output=processor_outputs) - true_length += image_offsets + true_length += image_offsets # pyrefly: ignore[unbound-name] if config.use_mrope: from maxtext.multimodal import processor_qwen3_omni # pylint: disable=import-outside-toplevel @@ -181,13 +181,13 @@ def main(argv: Sequence[str]) -> None: prefill_result, first_token = engine.prefill( params=params, padded_tokens=tokens, - positions=position_ids, - mrope_deltas=mrope_position_deltas, - images=processor_outputs.pixel_values if config.use_multimodal else None, - image_masks=processor_outputs.pixel_mask if config.use_multimodal and "llama4" in config.model_name else None, + positions=position_ids, # pyrefly: ignore[bad-argument-type] + mrope_deltas=mrope_position_deltas, # pyrefly: ignore[bad-argument-type] + images=processor_outputs.pixel_values if config.use_multimodal else None, # pyrefly: ignore[bad-argument-type] + image_masks=processor_outputs.pixel_mask if config.use_multimodal and "llama4" in config.model_name else None, # pyrefly: ignore[bad-argument-type] videos=getattr(processor_outputs, "video_values", None) if config.use_multimodal else None, - audio_values=processor_outputs.audio_values if config.use_audio else None, - audio_masks=processor_outputs.audio_mask if config.use_audio else None, + audio_values=processor_outputs.audio_values if config.use_audio else None, # pyrefly: ignore[bad-argument-type] + audio_masks=processor_outputs.audio_mask if config.use_audio else None, # pyrefly: ignore[bad-argument-type] true_length=true_length, rng=rng_prefill, slot=i, diff --git a/src/maxtext/inference/inference_microbenchmark.py b/src/maxtext/inference/inference_microbenchmark.py index 1adb0f36d5..d129883af8 100644 --- a/src/maxtext/inference/inference_microbenchmark.py +++ b/src/maxtext/inference/inference_microbenchmark.py @@ -400,8 +400,8 @@ def run_benchmarks(config): config, multisampling_prefill_executable[prefill_length], params, - prefill_tokens[prefill_length], - prefill_true_lengths[prefill_length], + prefill_tokens[prefill_length], # pyrefly: ignore[unbound-name] + prefill_true_lengths[prefill_length], # pyrefly: ignore[unbound-name] benchmark_loop_iters, ) diff --git a/src/maxtext/inference/kvcache.py b/src/maxtext/inference/kvcache.py index ba2266060f..52a3706028 100644 --- a/src/maxtext/inference/kvcache.py +++ b/src/maxtext/inference/kvcache.py @@ -277,7 +277,7 @@ def __init__( *, # Not used in KVCache but passed in by nnx_wrappers.to_linen. # TODO: Remove when bridge no longer needed - rngs: nnx.Rngs = None, + rngs: nnx.Rngs = None, # pyrefly: ignore[bad-function-definition] ): """Initializes the KVCache module. @@ -612,8 +612,8 @@ def kv_cache_chunked_prefill( # For quantized kv cached. Could be get without transpose twice. cached_key = self.get_cached_values(cached_prefill_key_vars, key.dtype, self.prefill_cache_axis_order) cached_value = self.get_cached_values(cached_prefill_value_vars, value.dtype, self.prefill_cache_axis_order) - cached_key_value = jnp.transpose(cached_key, self.prefill_cache_axis_order) - cached_value_value = jnp.transpose(cached_value, self.prefill_cache_axis_order) + cached_key_value = jnp.transpose(cached_key, self.prefill_cache_axis_order) # pyrefly: ignore[bad-argument-type] + cached_value_value = jnp.transpose(cached_value, self.prefill_cache_axis_order) # pyrefly: ignore[bad-argument-type] seq_axis = self.prefill_cache_logical_axis_names.index(CACHE_SEQUENCE) cache_seq_axis = self.prefill_cache_axis_order.index(seq_axis) @@ -767,17 +767,17 @@ def update_ar_key_value( if use_ragged_attention: cache_locations = [slice(None)] * 4 new_token_locations = [slice(None)] * 4 - new_token_locations[ar_cache_sequence_axis] = 0 + new_token_locations[ar_cache_sequence_axis] = 0 # pyrefly: ignore[unsupported-operation] def key_body(i, val): cache_locations[ar_cache_batch_axis] = i - cache_locations[ar_cache_sequence_axis] = lengths[i] + cache_locations[ar_cache_sequence_axis] = lengths[i] # pyrefly: ignore[unsupported-operation] new_token_locations[ar_cache_batch_axis] = i return val.at[tuple(cache_locations)].set(one_token_key_shaped_for_cache[tuple(new_token_locations)]) def value_body(i, val): cache_locations[ar_cache_batch_axis] = i - cache_locations[ar_cache_sequence_axis] = lengths[i] + cache_locations[ar_cache_sequence_axis] = lengths[i] # pyrefly: ignore[unsupported-operation] new_token_locations[ar_cache_batch_axis] = i return val.at[tuple(cache_locations)].set(one_token_value_shaped_for_cache[tuple(new_token_locations)]) @@ -815,7 +815,7 @@ def value_body(i, val): cached_key_scale.set_value( jax.lax.dynamic_update_index_in_dim( cached_key_scale.get_value(), - one_token_key_scale_shaped_for_cache, + one_token_key_scale_shaped_for_cache, # pyrefly: ignore[unbound-name] ar_cache_update_idx, ar_cache_scale_update_axis, ) @@ -823,7 +823,7 @@ def value_body(i, val): cached_value_scale.set_value( jax.lax.dynamic_update_index_in_dim( cached_value_scale.get_value(), - one_token_value_scale_shaped_for_cache, + one_token_value_scale_shaped_for_cache, # pyrefly: ignore[unbound-name] ar_cache_update_idx, ar_cache_scale_update_axis, ) @@ -844,7 +844,7 @@ def get_cached_values(self, cache_vars, target_dtype, cache_axis_order) -> jax.A elif dtype == jnp.float8_e4m3fn: scale_value /= E4M3_MAX - cache_value = KVTensor(qvalue=cache_value, scale=[scale_value], scale_t=None, dequant_dtype=target_dtype, bias=[]) + cache_value = KVTensor(qvalue=cache_value, scale=[scale_value], scale_t=None, dequant_dtype=target_dtype, bias=[]) # pyrefly: ignore[unexpected-keyword] cache_value_in_logical_shape = jax.tree.map(lambda x: reverse_transpose(x, cache_axis_order), cache_value) return cache_value_in_logical_shape @@ -1061,7 +1061,7 @@ def __init__( *, # Not used in MlaKVCache but passed in by nnx_wrappers.to_linen. # TODO: Remove when bridge no longer needed - rngs: nnx.Rngs = None, + rngs: nnx.Rngs = None, # pyrefly: ignore[bad-function-definition] ): """Initializes the MlaKVCache module. diff --git a/src/maxtext/inference/maxengine/maxengine.py b/src/maxtext/inference/maxengine/maxengine.py index 5377e86ec3..095d3ea48a 100644 --- a/src/maxtext/inference/maxengine/maxengine.py +++ b/src/maxtext/inference/maxengine/maxengine.py @@ -101,7 +101,7 @@ def get_keys(self): _BaseEngine = engine_api.Engine if (not is_decoupled() and hasattr(engine_api, "Engine")) else object -class MaxEngine(_BaseEngine): +class MaxEngine(_BaseEngine): # pyrefly: ignore[invalid-inheritance] """The computational core of the generative model server. Engine defines an API that models must adhere to as they plug into the @@ -192,7 +192,7 @@ def _nnx_init_cache_dict(self, mode: str = MODEL_MODE_PREFILL) -> dict: """Zero-filled pure-dict cache matching the abstract NNX model.""" src = self._abstract_model_for_mode(mode) _, cache_state, _ = nnx.split(src, nnx.Cache, ...) - cache_dict = cache_state.to_pure_dict() + cache_dict = cache_state.to_pure_dict() # pyrefly: ignore[missing-attribute] return jax.tree.map(lambda x: jnp.zeros(x.shape, x.dtype), cache_dict) def _nnx_run_model( @@ -219,7 +219,7 @@ def _nnx_run_model( nnx.replace_by_pure_dict(cache_state, cache_dict) # copy=True avoids reusing Variable objects across traces (TraceContextError), # mirroring the workaround in train.py's diff_wrapper. - model = nnx.merge(self.graphdef, params, cache_state, self._nnx_rest_state, copy=True) + model = nnx.merge(self.graphdef, params, cache_state, self._nnx_rest_state, copy=True) # pyrefly: ignore[no-matching-overload] logits = model( decoder_input_tokens, decoder_positions, @@ -280,7 +280,7 @@ def _layout(x, s, l): if x.format == l: return x # Somehow this can be None sometimes. - dll = (l.layout if jax.__version_info__ >= (0, 6, 3) else l.device_local_layout) if isinstance(l, Format) else l + dll = (l.layout if jax.__version_info__ >= (0, 6, 3) else l.device_local_layout) if isinstance(l, Format) else l # pyrefly: ignore[missing-attribute] f = jax.jit(self._identity, out_shardings=Format(dll, s)).lower(x).compile(compiler_options=xla_flags) y = f(x) # Achieves donation of the input argument, but allows for different memory @@ -409,7 +409,7 @@ def _load_params_nnx(self, params, rng): with nn_partitioning.axis_rules(self.config.logical_axis_rules): full_sharding = sharding.nnx_construct_named_sharding(full_abs, self._mesh) concrete_model = maxtext_utils_nnx.create_nnx_sharded_model( - self.model, self._create_model_fn, mesh=self._mesh, named_sharding=full_sharding + self.model, self._create_model_fn, mesh=self._mesh, named_sharding=full_sharding # pyrefly: ignore[bad-argument-type] ) graphdef, _, _, rest_state = nnx.split(concrete_model, nnx.Param, nnx.Cache, ...) self.graphdef = graphdef @@ -435,7 +435,7 @@ def _load_params_nnx(self, params, rng): # PREFILL/AR attention ops have different cache variable shapes, and a # mismatch trips the `assert prefill_kv_cache` check inside attention_op. with nn_partitioning.axis_rules(self.config.logical_axis_rules): - concrete_model = self._create_model_fn() + concrete_model = self._create_model_fn() # pyrefly: ignore[not-callable] graphdef, _, _, rest_state = nnx.split(concrete_model, nnx.Param, nnx.Cache, ...) # Overlay loaded non-Param/non-Cache leaves (e.g. AQT qrhs.frozen) onto # the PREFILL-mode rest_state. The PREFILL concrete_model already has @@ -443,8 +443,8 @@ def _load_params_nnx(self, params, rng): # values. Anything only in `loaded_rest_state` (e.g. AR-only RNG slots) # is ignored. We keep PREFILL rest_state as the base so RNG variables # match the PREFILL graphdef's expectations. - loaded_rest_dict = loaded_rest_state.to_pure_dict() - rest_dict = rest_state.to_pure_dict() + loaded_rest_dict = loaded_rest_state.to_pure_dict() # pyrefly: ignore[missing-attribute] + rest_dict = rest_state.to_pure_dict() # pyrefly: ignore[missing-attribute] def _overlay(dst, src): if isinstance(dst, dict) and isinstance(src, dict): @@ -768,7 +768,7 @@ def _prefill_jit( one_d_output = ones_to_keep * DECODING_ACTIVE_SEQUENCE_INDICATOR sequence_indicator = jnp.expand_dims(one_d_output, 0) - rng, new_rng = jax.random.split(rng) + rng, new_rng = jax.random.split(rng) # pyrefly: ignore[bad-argument-type] if self.config.pure_nnx: # Prefill always operates on batch=1 (one padded prompt at a time). nnx_cache = ( @@ -1026,7 +1026,7 @@ def _prefill_multisampling_jit( one_d_output = ones_to_keep * DECODING_ACTIVE_SEQUENCE_INDICATOR sequence_indicator = jnp.expand_dims(one_d_output, 0) - rng, new_rng = jax.random.split(rng) + rng, new_rng = jax.random.split(rng) # pyrefly: ignore[bad-argument-type] if self.config.pure_nnx: # Prefill is batch=1 (one prompt); multi-sampling only draws several first # tokens from the shared logits below. Mirror the _prefill_jit NNX branch. diff --git a/src/maxtext/inference/vllm_decode.py b/src/maxtext/inference/vllm_decode.py index d54d40c2d9..4cb6571686 100644 --- a/src/maxtext/inference/vllm_decode.py +++ b/src/maxtext/inference/vllm_decode.py @@ -36,26 +36,25 @@ from absl import flags from flax.linen import partitioning as nn_partitioning import jax -import transformers - -from maxtext.utils import model_creation_utils -from maxtext.utils import max_logging -from maxtext.utils.globals import MAXTEXT_CONFIGS_DIR -from maxtext.common.common_types import Config from maxtext.common import profiler +from maxtext.common.common_types import Config +from maxtext.configs import pyconfig from maxtext.integration.tunix.tunix_adapter import TunixMaxTextAdapter +import maxtext.integration.vllm.maxtext_vllm_adapter as adapter +from maxtext.utils import lora_utils +from maxtext.utils import max_logging +from maxtext.utils import model_creation_utils +from maxtext.utils.globals import MAXTEXT_CONFIGS_DIR +import transformers from tunix.rl.rollout import base_rollout from tunix.rl.rollout.vllm_rollout import VllmRollout from vllm import LLM -from vllm.sampling_params import SamplingParams -from maxtext.configs import pyconfig -import maxtext.integration.vllm.maxtext_vllm_adapter as adapter - # Force uses_mrope to False to disable 3D multimodal position IDs in text-only runs. # TODO(b/520142315): Class-level monkey patching is required here because ModelConfig.uses_mrope # is evaluated during the internal initialization of the LLM engine, making instance-level # configuration impossible. Check if a cleaner configuration option can be added in upstream vLLM. from vllm.config import ModelConfig +from vllm.sampling_params import SamplingParams ModelConfig.uses_mrope = property(lambda _: False) @@ -94,11 +93,13 @@ def decode_with_vllm(config: Config) -> None: "additional_config": { "maxtext_config": { "model_name": config.model_name, - "weight_dtype": "bfloat16", + "weight_dtype": config.weight_dtype.name, "allow_split_physical_axes": True, "debug_sharding": config.debug_sharding, "prefuse_moe_weights": config.prefuse_moe_weights, "scan_layers": config.scan_layers, + "enable_nnx": config.enable_nnx, + "pure_nnx_decoder": config.pure_nnx_decoder, }, "sharding": { "sharding_strategy": { @@ -195,7 +196,18 @@ def decode_with_tunix( mesh: The JAX mesh for parallelism. """ # Wrap the model for Tunix - tunix_model = TunixMaxTextAdapter(base_model=model) + use_no_op_mappings = False + if hasattr(config, "vllm_hf_overrides") and config.vllm_hf_overrides: + overrides = config.vllm_hf_overrides + if isinstance(overrides, str) and "MaxTextForCausalLM" in overrides: + use_no_op_mappings = True + elif isinstance(overrides, dict) and "MaxTextForCausalLM" in overrides.get("architectures", []): + use_no_op_mappings = True + + tunix_model = TunixMaxTextAdapter( + base_model=model, + use_no_op_mappings=use_no_op_mappings, + ) # Load the tokenizer tokenizer = transformers.AutoTokenizer.from_pretrained( @@ -223,13 +235,48 @@ def decode_with_tunix( f"max_target_length ({config.max_target_length}) must be greater than max_prompt_length ({max_prompt_length})" ) + # MaxText uses -1 to mean "disabled"; vLLM requires top_p in (0, 1]. + top_p = config.decode_sampling_nucleus_p if config.decode_sampling_nucleus_p > 0 else 1.0 + top_k = config.decode_sampling_top_k if config.decode_sampling_top_k > 0 else -1 + + rollout_vllm_additional_config = { + "maxtext_config": { + "model_name": config.model_name, + "weight_dtype": config.weight_dtype.name, + "allow_split_physical_axes": True, + "debug_sharding": config.debug_sharding, + "prefuse_moe_weights": config.prefuse_moe_weights, + "scan_layers": config.scan_layers, + "enable_nnx": config.enable_nnx, + "pure_nnx_decoder": config.pure_nnx_decoder, + } + } + + if config.lora.enable_lora: + rollout_vllm_additional_config["maxtext_config"]["lora"] = { + "enable_lora": config.lora.enable_lora, + "lora_restore_path": config.lora.lora_restore_path, + "lora_rank": config.lora.lora_rank, + "lora_alpha": config.lora.lora_alpha, + "lora_module_path": config.lora.lora_module_path, + } + # Create vLLM rollout for inference rollout_config = base_rollout.RolloutConfig( max_tokens_to_generate=max_tokens_to_generate, max_prompt_length=max_prompt_length, temperature=config.decode_sampling_temperature, - top_p=config.decode_sampling_nucleus_p, - top_k=config.decode_sampling_top_k, + top_p=top_p, + top_k=top_k, + rollout_vllm_model_version=config.tokenizer_path, + rollout_vllm_hbm_utilization=config.hbm_utilization_vllm, + rollout_vllm_init_with_random_weights=True, + rollout_vllm_tpu_backend_type="jax", + tensor_parallel_size=(config.ici_tensor_parallelism if config.ici_tensor_parallelism > 0 else 1), + data_parallel_size=jax.device_count() + // (config.ici_tensor_parallelism if config.ici_tensor_parallelism > 0 else 1), + rollout_vllm_additional_config=rollout_vllm_additional_config, + rollout_vllm_kwargs={"hf_overrides": config.vllm_hf_overrides}, ) vllm_rollout = VllmRollout( model=tunix_model, @@ -239,12 +286,7 @@ def decode_with_tunix( # other special formatting, which is not part of max_prompt_length. cache_config_or_size=max_prompt_length + max_tokens_to_generate + 256, mesh=mesh, - model_version=config.tokenizer_path, - hbm_utilization=0.8, - # Initialize vllm model with random weights to speed up bootstrap time. - # Actual model weights will be loaded later. - init_with_random_weights=True, - tpu_backend_type="jax", + rollout_config=rollout_config, ) # Generate text @@ -271,6 +313,10 @@ def main(argv: Sequence[str]) -> None: if FLAGS.use_tunix: maxtext_model, mesh = model_creation_utils.from_pretrained(config) + if config.lora.enable_lora: + maxtext_model = lora_utils.apply_lora_to_model(maxtext_model, mesh, config) + if config.lora.lora_restore_path: + lora_utils.restore_lora_from_path(maxtext_model, config) decode_with_tunix(config, model=maxtext_model, mesh=mesh) else: decode_with_vllm(config) diff --git a/src/maxtext/input_pipeline/grain_data_processing.py b/src/maxtext/input_pipeline/grain_data_processing.py index 9141f3bb1c..220c3ce82d 100644 --- a/src/maxtext/input_pipeline/grain_data_processing.py +++ b/src/maxtext/input_pipeline/grain_data_processing.py @@ -191,12 +191,11 @@ def create_dataset_from_pattern(pattern): f"Each shard will be read by at most {math.ceil(dataloading_host_count / len(data_files))} hosts. " f"Concurrent reading by multiple hosts may cause slow down." ) - min_files_per_host = len(data_files) // dataloading_host_count - if grain_worker_count > min_files_per_host: + if grain_worker_count > files_per_host: raise ValueError( - f"grain_worker_count ({grain_worker_count}) exceeds the minimum number of {data_file_type} files " - f"per host ({min_files_per_host} = {len(data_files)} files / {dataloading_host_count} hosts). " - f"Lower grain_worker_count to at most {min_files_per_host}." + f"grain_worker_count ({grain_worker_count}) exceeds the number of {data_file_type} files " + f"per host ({files_per_host}). " + f"Lower grain_worker_count to at most {files_per_host}." ) dataset = grain.MapDataset.source(data_files) if shuffle: @@ -204,9 +203,9 @@ def create_dataset_from_pattern(pattern): dataset = dataset.repeat(num_epoch) dataset = dataset[file_slice] if data_file_type == "tfrecord": - dataset = dataset.map(input_pipeline_utils.make_tfrecord_iter_dataset) + dataset = dataset.map(input_pipeline_utils.make_tfrecord_iter_dataset) # pyrefly: ignore[missing-attribute] else: - dataset = dataset.map(grain.experimental.ParquetIterDataset) + dataset = dataset.map(grain.experimental.ParquetIterDataset) # pyrefly: ignore[missing-attribute] cycle_length = min(files_per_host, grain_num_threads) dataset = grain.experimental.InterleaveIterDataset(dataset, cycle_length=cycle_length) if row_shard is not None: diff --git a/src/maxtext/input_pipeline/grain_tokenizer.py b/src/maxtext/input_pipeline/grain_tokenizer.py index 0ca4d222db..06556f4027 100644 --- a/src/maxtext/input_pipeline/grain_tokenizer.py +++ b/src/maxtext/input_pipeline/grain_tokenizer.py @@ -71,7 +71,7 @@ class TokenizeAndTrim(TokenizerTransformBase, grain.MapTransform): def map(self, element: dict[str, Any]) -> dict[str, Any]: """Maps to each element.""" - for feature_name, max_length in zip(self.feature_names, self.sequence_length, strict=True): + for feature_name, max_length in zip(self.feature_names, self.sequence_length, strict=True): # pyrefly: ignore[bad-argument-type] text = element[feature_name] token_ids = self._encode(text)[:max_length] element[feature_name] = np.asarray(token_ids, dtype=np.int32) @@ -88,9 +88,9 @@ def __post_init__(self): super().__post_init__() # TokenizeAndChunk only supports single feature for chunking assert len(self.feature_names) == 1, "TokenizeAndChunk only supports single feature name" - assert len(self.sequence_length) == 1, "TokenizeAndChunk only supports single sequence length" + assert len(self.sequence_length) == 1, "TokenizeAndChunk only supports single sequence length" # pyrefly: ignore[bad-argument-type] self.feature_name = self.feature_names[0] # For backward compatibility - self.sequence_length = self.sequence_length[0] # Convert back to int for chunking + self.sequence_length = self.sequence_length[0] # Convert back to int for chunking # pyrefly: ignore[bad-index] def flat_map(self, element: dict[str, Any]) -> list[dict[str, Any]]: """Tokenize and chunk text into multiple examples of sequence length.""" @@ -103,8 +103,8 @@ def flat_map(self, element: dict[str, Any]) -> list[dict[str, Any]]: return [] output_elements = [] - for start_idx in range(0, len(token_ids), chunk_size): - chunk = np.asarray(token_ids[start_idx : start_idx + chunk_size], dtype=np.int32) + for start_idx in range(0, len(token_ids), chunk_size): # pyrefly: ignore[bad-argument-type] + chunk = np.asarray(token_ids[start_idx : start_idx + chunk_size], dtype=np.int32) # pyrefly: ignore[unsupported-operation] new_element = {self.feature_name: chunk} output_elements.append(new_element) diff --git a/src/maxtext/input_pipeline/input_pipeline_utils.py b/src/maxtext/input_pipeline/input_pipeline_utils.py index 8ba57363fb..0b7751b522 100644 --- a/src/maxtext/input_pipeline/input_pipeline_utils.py +++ b/src/maxtext/input_pipeline/input_pipeline_utils.py @@ -172,7 +172,7 @@ def is_conversational(features, data_columns): for column in data_columns: messages = features[column] if isinstance(messages, datasets.Sequence): - if isinstance(messages.feature, dict) and "role" in messages.feature and "content" in messages.feature: + if isinstance(messages.feature, dict) and "role" in messages.feature and "content" in messages.feature: # pyrefly: ignore[missing-attribute] return True return False @@ -778,7 +778,7 @@ def _pad_image_and_mask(self, preprocessed_image: mm_utils.PreprocessorOutput) - if preprocessed_image.pixel_values is None: raise ValueError("Input preprocessed_image must have pixel_values to pad images.") - if self.config.model_name and self.config.model_name.startswith("qwen3-omni"): + if self.config.model_name and self.config.model_name.startswith("qwen3-omni"): # pyrefly: ignore[missing-attribute] return preprocessed_image # Determine the maximum number of images/masks allowed. @@ -840,21 +840,21 @@ def map( if isinstance(element[data_column], mm_utils.PreprocessorOutput): raise TypeError("Only 'images' column can be of type PreprocessorOutput.") - element[f"{data_column}_segmentation"] = element[data_column] != self.pad_id - element[f"{data_column}_segmentation"] = element[f"{data_column}_segmentation"].astype(np.int32) - element[f"{data_column}_position"] = np.arange(element[data_column].shape[0], dtype=np.int32) + element[f"{data_column}_segmentation"] = element[data_column] != self.pad_id # pyrefly: ignore[unsupported-operation] + element[f"{data_column}_segmentation"] = element[f"{data_column}_segmentation"].astype(np.int32) # pyrefly: ignore[missing-attribute] + element[f"{data_column}_position"] = np.arange(element[data_column].shape[0], dtype=np.int32) # pyrefly: ignore[missing-attribute] if self.add_true_length: - element[f"{data_column}_true_length"] = np.array([element[data_column].shape[0]], dtype=np.int32) + element[f"{data_column}_true_length"] = np.array([element[data_column].shape[0]], dtype=np.int32) # pyrefly: ignore[missing-attribute] for key, _ in element.items(): if key == "images": - if self.config.model_name is None: + if self.config.model_name is None: # pyrefly: ignore[missing-attribute] raise ValueError("model_name must be provided when padding images") - element["images"] = self._pad_image_and_mask(element["images"]) + element["images"] = self._pad_image_and_mask(element["images"]) # pyrefly: ignore[bad-argument-type] elif "true_length" not in key: - element[key] = self._pad_text(element[key], self.max_length, self.pad_id) + element[key] = self._pad_text(element[key], self.max_length, self.pad_id) # pyrefly: ignore[bad-argument-type] return element @@ -881,7 +881,7 @@ def map(self, element: dict[str, np.ndarray]) -> dict[str, np.ndarray]: raise TypeError(f"'images' must be of type PreprocessorOutput, but got {type(preprocessed_image)}") output = element.copy() - output["images"] = preprocessed_image.pixel_values + output["images"] = preprocessed_image.pixel_values # pyrefly: ignore[unsupported-operation] if preprocessed_image.pixel_mask is not None: output["image_masks"] = preprocessed_image.pixel_mask diff --git a/src/maxtext/input_pipeline/multihost_dataloading.py b/src/maxtext/input_pipeline/multihost_dataloading.py index 41f40a7608..8e1246a25b 100644 --- a/src/maxtext/input_pipeline/multihost_dataloading.py +++ b/src/maxtext/input_pipeline/multihost_dataloading.py @@ -239,11 +239,11 @@ def save_state(self, step_array): if jax.process_index() == 0: directory.mkdir(parents=True, exist_ok=True) filename = directory / "process_0.json" - filename.write_text(json.dumps(self.iterator.get_state(), indent=4)) + filename.write_text(json.dumps(self.iterator.get_state(), indent=4)) # pyrefly: ignore[missing-attribute] return step_array directory.mkdir(parents=True, exist_ok=True) filename = directory / f"process_{jax.process_index()}-of-{jax.process_count()}.json" - state = json.dumps(self.iterator.get_state(), indent=4) + state = json.dumps(self.iterator.get_state(), indent=4) # pyrefly: ignore[missing-attribute] filename.write_text(state) return step_array @@ -255,7 +255,7 @@ def restore_state(self, step_array): else: filename = directory / f"process_{jax.process_index()}-of-{jax.process_count()}.json" state = json.loads(filename.read_text()) - self.iterator.set_state(state) + self.iterator.set_state(state) # pyrefly: ignore[missing-attribute] return step_array @@ -273,11 +273,11 @@ def __init__(self, get_ds_fn, preprocessing_fn, global_mesh, global_shape, check # named "local_iterator" to match MultiHostDataLoadIterator's interface. remote_iterator_cls = colocated_python.colocated_python_class(RemoteIterator) self.local_iterator = remote_iterator_cls( - get_ds_fn, + get_ds_fn, # pyrefly: ignore[bad-argument-count] preprocessing_fn, global_shape, checkpoint_path, - elastic=elastic, + elastic=elastic, # pyrefly: ignore[unexpected-keyword] ) max_logging.log("RemoteIteratorWrapper initiated") @@ -285,10 +285,10 @@ def __iter__(self): return self def reset(self): - self.local_iterator.reset() + self.local_iterator.reset() # pyrefly: ignore[missing-attribute] def __next__(self): - out = self.local_iterator.get_next(self.dummy_array) + out = self.local_iterator.get_next(self.dummy_array) # pyrefly: ignore[missing-attribute] # use tree_map is out is a dict return jax.device_put(out, self.tpu_sharding) @@ -296,10 +296,10 @@ def save_state(self, step): replicated_cpu_sharding = NamedSharding(self.cpu_mesh, PartitionSpec()) step_array = jnp.array(step, dtype=jnp.int32) step_array = jax.device_put(step_array, replicated_cpu_sharding) - self.local_iterator.save_state(step_array) + self.local_iterator.save_state(step_array) # pyrefly: ignore[missing-attribute] def restore_state(self, step): replicated_cpu_sharding = NamedSharding(self.cpu_mesh, PartitionSpec()) step_array = jnp.array(step, dtype=jnp.int32) step_array = jax.device_put(step_array, replicated_cpu_sharding) - self.local_iterator.restore_state(step_array) + self.local_iterator.restore_state(step_array) # pyrefly: ignore[missing-attribute] diff --git a/src/maxtext/input_pipeline/olmo_data.py b/src/maxtext/input_pipeline/olmo_data.py index 82c258b78a..39d5cdd031 100644 --- a/src/maxtext/input_pipeline/olmo_data.py +++ b/src/maxtext/input_pipeline/olmo_data.py @@ -137,7 +137,7 @@ def global_to_local(index: OlmoNpyIndex, instance_id: int) -> Tuple[int, int]: if instance_id < 0 or instance_id >= index.total_instances: raise IndexError(f"instance_id {instance_id} out of range " f"[0, {index.total_instances})") starts = index._instance_offset_starts # type: ignore[attr-defined] # pylint: disable=protected-access - file_idx = bisect.bisect_right(starts, instance_id) - 1 + file_idx = bisect.bisect_right(starts, instance_id) - 1 # pyrefly: ignore[bad-argument-type] local_instance = instance_id - index.files[file_idx].instance_offset token_offset = local_instance * index.sequence_length return file_idx, token_offset diff --git a/src/maxtext/input_pipeline/packing/prefill_packing.py b/src/maxtext/input_pipeline/packing/prefill_packing.py index 8cc8816437..d3f74f5ab2 100644 --- a/src/maxtext/input_pipeline/packing/prefill_packing.py +++ b/src/maxtext/input_pipeline/packing/prefill_packing.py @@ -330,7 +330,7 @@ def zero_padded(arr: list[int], padding: int): prompt_logp_numpy = np.array(decode_state["prompt_logp"]) for i in range(bucket.count): if return_prompt_logp: - prompt_logp = prompt_logp_numpy[:, offsets[i] : offsets[i] + lengths[i]] + prompt_logp = prompt_logp_numpy[:, offsets[i] : offsets[i] + lengths[i]] # pyrefly: ignore[unsupported-operation] prefill_result.append(PrefillResult(first_tokens[i], bucket.slots[i], prompt_logp)) else: prefill_result.append(PrefillResult(first_tokens[i], bucket.slots[i], None)) diff --git a/src/maxtext/integration/tunix/utils.py b/src/maxtext/integration/tunix/utils.py index 8d608956bc..5c8b9ba370 100644 --- a/src/maxtext/integration/tunix/utils.py +++ b/src/maxtext/integration/tunix/utils.py @@ -153,10 +153,60 @@ def to_hf_hook_fns(self): return {} def lora_to_hf_mappings(self): - if self.use_standalone_mappings: - return STANDALONE_VLLM_WEIGHT_MAPPING[self.model_name].lora_to_hf_mappings() + """Dynamically generate LoRA mappings from base model weights mappings.""" + base_mappings = self.to_hf_mapping() + if not base_mappings: + return None + + lora_mapping = {} + for maxtext_key, (hf_key, sharding_spec) in base_mappings.items(): + segments = set(maxtext_key.split(".")) + is_input_proj = any( + p in segments + for p in [ + "wi_0", + "wi_1", + "query", + "key", + "value", + "wq_a", + "wq_b", + "wkv_a", + "wkv_b", + ] + ) + is_output_proj = any(p in segments for p in ["wo", "out"]) + + if not (is_input_proj or is_output_proj): + continue + + # Derive MaxText LoRA keys + maxtext_lora_a = maxtext_key + "_lora_a" + maxtext_lora_b = maxtext_key + "_lora_b" + + # Derive HF/vLLM LoRA keys + if hf_key.endswith(".kernel"): + hf_lora_a = hf_key.replace(".kernel", ".kernel_lora_a") + hf_lora_b = hf_key.replace(".kernel", ".kernel_lora_b") + elif hf_key.endswith(".weight"): + hf_lora_a = hf_key.replace(".weight", ".weight_lora_a") + hf_lora_b = hf_key.replace(".weight", ".weight_lora_b") + else: + hf_lora_a = hf_key + "_lora_a" + hf_lora_b = hf_key + "_lora_b" + + # Derive sharding specifications for Qwix LoRA parameters + if is_input_proj: + sharding_a = (None, "layer", None) # Input -> Rank (unsharded) + sharding_b = sharding_spec # Rank -> Output (same as base) + else: + sharding_a = sharding_spec # Input -> Rank (same as base) + sharding_b = (None, "layer", None) # Rank -> Output (unsharded) + + lora_mapping[maxtext_lora_a] = (hf_lora_a, sharding_a) + lora_mapping[maxtext_lora_b] = (hf_lora_b, sharding_b) - return None + return lora_mapping def _generalize_maxtext_key(self, maxtext_key): """Generalizes the MaxText key to a common vLLM format.""" diff --git a/src/maxtext/integration/tunix/weight_mapping/__init__.py b/src/maxtext/integration/tunix/weight_mapping/__init__.py index 6fd7f35028..39ab12ff8f 100644 --- a/src/maxtext/integration/tunix/weight_mapping/__init__.py +++ b/src/maxtext/integration/tunix/weight_mapping/__init__.py @@ -19,6 +19,7 @@ model name. This allows for easy extension to support new models. """ from maxtext.integration.tunix.weight_mapping.deepseek3 import DEEPSEEK_VLLM_MAPPING +from maxtext.integration.tunix.weight_mapping.gemma3 import GEMMA3_VLLM_MAPPING from maxtext.integration.tunix.weight_mapping.gpt_oss import GPT_OSS_VLLM_MAPPING from maxtext.integration.tunix.weight_mapping.llama3 import LLAMA3_VLLM_MAPPING from maxtext.integration.tunix.weight_mapping.qwen2 import QWEN2_VLLM_MAPPING @@ -35,6 +36,8 @@ def __getattr__(self, name): return QWEN2_VLLM_MAPPING elif name.startswith("qwen3"): return QWEN3_VLLM_MAPPING + elif name.startswith("gemma3"): + return GEMMA3_VLLM_MAPPING elif name.startswith("deepseek3"): return DEEPSEEK_VLLM_MAPPING elif name.startswith("gpt-oss"): diff --git a/src/maxtext/integration/tunix/weight_mapping/gemma3.py b/src/maxtext/integration/tunix/weight_mapping/gemma3.py new file mode 100644 index 0000000000..de0bb47ed9 --- /dev/null +++ b/src/maxtext/integration/tunix/weight_mapping/gemma3.py @@ -0,0 +1,124 @@ +# Copyright 2023–2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Defines the weight mapping from MaxText's Gemma3 model to a vLLM-compatible format.""" + +import dataclasses + +import numpy as np + + +@dataclasses.dataclass +# pylint: disable=invalid-name +class GEMMA3_VLLM_MAPPING: + """Mapping MaxText Gemma3 weights to vLLM's Gemma3 weights.""" + + @staticmethod + def to_hf_hook_fns(): + """Returns a dictionary of hook functions to be applied to MaxText weights.""" + + def scale_embedding(arr): + hidden_size = arr.shape[1] + normalizer = np.dtype(arr.dtype).type(hidden_size**0.5) + return arr / normalizer + + return { + "base.token_embedder.embedding": scale_embedding, + } + + @staticmethod + def to_hf_transpose_keys(): + """Returns a list of keys for weights that need to be transposed.""" + return {} + + @staticmethod + def lora_to_hf_mappings(): + """Provides the mapping for LoRA (Low-Rank Adaptation) weights.""" + return None + + @staticmethod + def to_hf_mapping(): + """Mapping from MaxText model to HuggingFace vLLM model. + + Returns: + A dictionary mapping MaxText parameter names to HuggingFace parameter + names and sharding. + """ + return { + # Token embeddings - shard vocab dimension + "base.token_embedder.embedding": ( + "model.language_model.embed_tokens.kernel", + ("model", None), + ), + # Final layer norm - no sharding needed + "base.decoder.decoder_norm.scale": ( + "model.language_model.norm.scale", + (None,), + ), + # Layer norms - no sharding needed + "base.decoder.layers.pre_self_attention_norm.scale": ( + "model.language_model.layers.*.input_layernorm.scale", + (None, "layer"), + ), + "base.decoder.layers.post_self_attention_norm.scale": ( + "model.language_model.layers.*.post_attention_layernorm.scale", + (None, "layer"), + ), + "base.decoder.layers.self_attention.query_norm.scale": ( + "model.language_model.layers.*.self_attn.q_norm.scale", + (None, "layer"), + ), + "base.decoder.layers.self_attention.key_norm.scale": ( + "model.language_model.layers.*.self_attn.k_norm.scale", + (None, "layer"), + ), + "base.decoder.layers.pre_ffw_norm.scale": ( + "model.language_model.layers.*.pre_feedforward_layernorm.scale", + (None, "layer"), + ), + "base.decoder.layers.post_ffw_norm.scale": ( + "model.language_model.layers.*.post_feedforward_layernorm.scale", + (None, "layer"), + ), + # MLP components - shard hidden dimensions + "base.decoder.layers.mlp.wi_0.kernel": ( + "model.language_model.layers.*.mlp.gate_proj.kernel", + (None, "layer", "model"), + ), + "base.decoder.layers.mlp.wi_1.kernel": ( + "model.language_model.layers.*.mlp.up_proj.kernel", + (None, "layer", "model"), + ), + "base.decoder.layers.mlp.wo.kernel": ( + "model.language_model.layers.*.mlp.down_proj.kernel", + ("model", "layer", None), + ), + # Attention components - shard head dimensions + "base.decoder.layers.self_attention.query.kernel": ( + "model.language_model.layers.*.self_attn.q_proj.kernel", + (None, "layer", "model", None), + ), + "base.decoder.layers.self_attention.key.kernel": ( + "model.language_model.layers.*.self_attn.k_proj.kernel", + (None, "layer", "model", None), + ), + "base.decoder.layers.self_attention.value.kernel": ( + "model.language_model.layers.*.self_attn.v_proj.kernel", + (None, "layer", "model", None), + ), + "base.decoder.layers.self_attention.out.kernel": ( + "model.language_model.layers.*.self_attn.o_proj.kernel", + ("model", "layer", None, None), + ), + } diff --git a/src/maxtext/integration/vllm/maxtext_vllm_adapter/adapter.py b/src/maxtext/integration/vllm/maxtext_vllm_adapter/adapter.py index e3f2785fde..1d441327ac 100644 --- a/src/maxtext/integration/vllm/maxtext_vllm_adapter/adapter.py +++ b/src/maxtext/integration/vllm/maxtext_vllm_adapter/adapter.py @@ -15,19 +15,18 @@ """vLLM adapter for MaxText models.""" import os -import jax - from flax import nnx import flax.linen as nn +import jax from jax import numpy as jnp from jax.experimental.pallas import tpu as pltpu from jax.sharding import Mesh - -from maxtext.configs import pyconfig -from maxtext.utils.globals import MAXTEXT_CONFIGS_DIR from maxtext.common.common_types import MODEL_MODE_AUTOREGRESSIVE +from maxtext.configs import pyconfig +from maxtext.utils import lora_utils from maxtext.utils import max_logging from maxtext.utils import model_creation_utils +from maxtext.utils.globals import MAXTEXT_CONFIGS_DIR try: @@ -345,6 +344,12 @@ def load_weights(self, rng_key: jax.Array) -> None: model = model_creation_utils.from_pretrained( self.maxtext_config, mesh=self.mesh, model_mode=self.model_mode, rng_key=rng_key ) + if self.maxtext_config.lora.enable_lora: + model = lora_utils.apply_lora_to_model( + model, self.mesh, self.maxtext_config + ) + if self.maxtext_config.lora.lora_restore_path: + lora_utils.restore_lora_from_path(model, self.maxtext_config) self.model = nnx.data(model) def get_mrope_input_positions( diff --git a/src/maxtext/kernels/attention/jax_flash_attention.py b/src/maxtext/kernels/attention/jax_flash_attention.py index 7b8b3b94b1..93a1933615 100644 --- a/src/maxtext/kernels/attention/jax_flash_attention.py +++ b/src/maxtext/kernels/attention/jax_flash_attention.py @@ -159,10 +159,10 @@ def compute_attention_block(output, l, m): (batch_size, num_kv_heads, q_groups, block_q, block_kv), ) - s_i_j = jnp.where(broadcasted_mask, s_i_j, mask_value) if cap is not None: s_i_j = jnp.tanh(s_i_j / cap) s_i_j = s_i_j * cap + s_i_j = jnp.where(broadcasted_mask, s_i_j, mask_value) m_i_j = s_i_j.max(axis=-1) p_i_j = jnp.exp(s_i_j - m_i_j[..., None]) l_i_j = p_i_j.sum(axis=-1) diff --git a/src/maxtext/kernels/attention/splash_attention_kernel.py b/src/maxtext/kernels/attention/splash_attention_kernel.py index bd5099972f..0eb57087f6 100644 --- a/src/maxtext/kernels/attention/splash_attention_kernel.py +++ b/src/maxtext/kernels/attention/splash_attention_kernel.py @@ -114,7 +114,7 @@ def get_kernel_name( @overload -def _attention_reference( +def _attention_reference( # pyrefly: ignore[inconsistent-overload] mask: jax.Array, q: jax.Array, k: jax.Array, @@ -130,7 +130,7 @@ def _attention_reference( @overload -def _attention_reference( +def _attention_reference( # pyrefly: ignore[inconsistent-overload] mask: jax.Array, q: jax.Array, k: jax.Array, @@ -618,18 +618,18 @@ def _apply_mask_and_soft_cap( # need to keep into account the current shard along Q sequence. if k_in_lanes: - assert q_sequence_ref.shape == (bq, NUM_LANES) + assert q_sequence_ref.shape == (bq, NUM_LANES) # pyrefly: ignore[missing-attribute] k_sequence = k_offset + jax.lax.broadcasted_iota(jnp.int32, (bq, k_slice.size), 1) repeats, rem = divmod(k_slice.size, NUM_LANES) assert rem == 0 - q_sequence = jnp.tile(q_sequence_ref[...], (1, repeats)) # [bq, k_slice.size] + q_sequence = jnp.tile(q_sequence_ref[...], (1, repeats)) # [bq, k_slice.size] # pyrefly: ignore[unsupported-operation] else: - assert q_sequence_ref.shape == (NUM_SUBLANES, bq) + assert q_sequence_ref.shape == (NUM_SUBLANES, bq) # pyrefly: ignore[missing-attribute] k_sequence = k_offset + jax.lax.broadcasted_iota(jnp.int32, (k_slice.size, bq), 0) - q_sequence = q_sequence_ref[:1, :] # [1, bq] + q_sequence = q_sequence_ref[:1, :] # [1, bq] # pyrefly: ignore[unsupported-operation] q_sequence = jnp.broadcast_to(q_sequence, (k_slice.size, bq)) assert q_sequence.shape == k_sequence.shape @@ -773,7 +773,7 @@ def body(kv_compute_index, _): if rem != 0: raise NotImplementedError(f"{bkv_compute=} should be a multiple of {NUM_LANES}") - s_curr = jnp.exp(qk - jnp.tile(m_next, (1, bkv_repeats))) + s_curr = jnp.exp(qk - jnp.tile(m_next, (1, bkv_repeats))) # pyrefly: ignore[unsupported-operation] assert s_curr.shape == (bq, bkv_compute) l_curr = jax.lax.broadcast_in_dim(s_curr.sum(axis=-1), l_prev.shape, (0,)) @@ -815,7 +815,7 @@ def end(): @overload -def _splash_attention_forward( +def _splash_attention_forward( # pyrefly: ignore[inconsistent-overload] fwd_mask_info: mask_info_lib.MaskInfo, q: jax.Array, k: jax.Array, @@ -833,7 +833,7 @@ def _splash_attention_forward( @overload -def _splash_attention_forward( +def _splash_attention_forward( # pyrefly: ignore[inconsistent-overload] fwd_mask_info: mask_info_lib.MaskInfo, q: jax.Array, k: jax.Array, @@ -908,9 +908,9 @@ def _splash_attention_forward( if k.shape[:-1] != v.shape[:-1]: raise ValueError(f"Expected 'key' {k.shape} and 'value' {v.shape} to have the same " "leading dimensions.") - if bkv % bkv_compute: + if bkv % bkv_compute: # pyrefly: ignore[unsupported-operation] raise ValueError(f"{bkv=} must be a multiple of {bkv_compute=}.") - if bkv_compute % NUM_LANES: + if bkv_compute % NUM_LANES: # pyrefly: ignore[unsupported-operation] raise ValueError(f"{bkv_compute=} must be a multiple of {NUM_LANES}.") kv_seq_len = k.shape[kv_seq_len_dimension] @@ -988,7 +988,7 @@ def kv_segment_ids_index_map(h, i, j, data_next_ref, block_mask_ref, mask_next_r if fwd_mask_info.partial_mask_blocks is not None: in_specs.append(pl.BlockSpec((None, bq, bkv), mask_index_map)) else: - in_specs.append(None) + in_specs.append(None) # pyrefly: ignore[bad-argument-type] assert fwd_mask_info.partial_mask_blocks is None or fwd_mask_info.q_sequence is None @@ -997,7 +997,7 @@ def kv_segment_ids_index_map(h, i, j, data_next_ref, block_mask_ref, mask_next_r in_specs.append(pl.BlockSpec((bq, NUM_LANES), q_segment_ids_index_map)) else: q_sequence = None - in_specs.append(None) + in_specs.append(None) # pyrefly: ignore[bad-argument-type] num_scalar_prefetch = 3 @@ -1266,8 +1266,8 @@ def run(): q_sequence_ref, q_segment_ids_ref, kv_segment_ids_ref, - attn_logits_soft_cap=attn_logits_soft_cap, - k_slice=pl.ds(0, bkv), + attn_logits_soft_cap=attn_logits_soft_cap, # pyrefly: ignore[bad-argument-type] + k_slice=pl.ds(0, bkv), # pyrefly: ignore[bad-argument-type] # When the iteration space is shrunk (for local attention for example), # the kv_index program_id does not correspond to the actual coordinates # of the KV data. Make sure to use the 'unshrunk' index (coming from the @@ -1276,7 +1276,7 @@ def run(): bq=bq, mask_function=mask_function, ) - p = jnp.exp(qk - logsumexp) + p = jnp.exp(qk - logsumexp) # pyrefly: ignore[unsupported-operation] dp_dims = NT_DIM_NUMBERS if v_layout == HEAD_DIM_MINOR else NN_DIM_NUMBERS dp = lax.dot_general( do.astype(v.dtype), @@ -1426,11 +1426,11 @@ def logsumexp_index_map(h, i, *_): logsumexp = jnp.expand_dims(logsumexp, axis=-2) logsumexp_spec = pl.BlockSpec((None, 1, bq), logsumexp_index_map) - assert logsumexp.ndim == len(logsumexp_spec.block_shape) + assert logsumexp.ndim == len(logsumexp_spec.block_shape) # pyrefly: ignore[bad-argument-type] di = jnp.expand_dims(di, axis=-2) di_spec = pl.BlockSpec((None, 1, bq), logsumexp_index_map) - assert di.ndim == len(di_spec.block_shape) + assert di.ndim == len(di_spec.block_shape) # pyrefly: ignore[bad-argument-type] in_specs = [ q_spec, @@ -1633,8 +1633,8 @@ def _load_kv(ref, layout): q_sequence_ref, q_segment_ids_ref, kv_segment_ids_ref, - attn_logits_soft_cap=attn_logits_soft_cap, - k_slice=slice_k, + attn_logits_soft_cap=attn_logits_soft_cap, # pyrefly: ignore[bad-argument-type] + k_slice=slice_k, # pyrefly: ignore[bad-argument-type] k_offset=kv_index * bkv + i * bkv_compute, bq=bq, k_in_lanes=False, @@ -1694,7 +1694,7 @@ def run(): if is_mqa: should_write = jnp.logical_and(should_write, q_head_index == num_q_heads - 1) elif num_kv_heads < num_q_heads: - should_write = jnp.logical_and(should_write, q_head_index_per_kv_head == q_heads_per_kv_heads - 1) + should_write = jnp.logical_and(should_write, q_head_index_per_kv_head == q_heads_per_kv_heads - 1) # pyrefly: ignore[unsupported-operation] @pl.when(should_write) def end(): @@ -1943,12 +1943,12 @@ def logsumexp_index_map( logsumexp_shape = (num_q_heads, NUM_SUBLANES, q_seq_len) logsumexp = jnp.broadcast_to(jnp.expand_dims(logsumexp, -2), logsumexp_shape) logsumexp_spec = pl.BlockSpec((None, NUM_SUBLANES, bq), logsumexp_index_map) - assert logsumexp.ndim == len(logsumexp_spec.block_shape) + assert logsumexp.ndim == len(logsumexp_spec.block_shape) # pyrefly: ignore[bad-argument-type] # TODO(apaszke): Remove the sublane expansion once Mosaic has all retilings di = jnp.broadcast_to(jnp.expand_dims(di, -2), logsumexp_shape) di_spec = pl.BlockSpec((None, NUM_SUBLANES, bq), logsumexp_index_map) - assert di.ndim == len(di_spec.block_shape) + assert di.ndim == len(di_spec.block_shape) # pyrefly: ignore[bad-argument-type] in_specs = [ q_spec, @@ -2117,11 +2117,11 @@ def _splash_attention_bwd( logsumexp, do, di, - bq=bq_dkv, - bkv=bkv_dkv_memory, - bkv_compute=bkv_dkv_compute, + bq=bq_dkv, # pyrefly: ignore[bad-argument-type] + bkv=bkv_dkv_memory, # pyrefly: ignore[bad-argument-type] + bkv_compute=bkv_dkv_compute, # pyrefly: ignore[bad-argument-type] is_mqa=is_mqa, - mask_info=dkv_mask_info, + mask_info=dkv_mask_info, # pyrefly: ignore[bad-argument-type] mask_value=mask_value, attn_logits_soft_cap=attn_logits_soft_cap, use_fused_bwd_kernel=use_fused_bwd_kernel, @@ -2141,10 +2141,10 @@ def _splash_attention_bwd( logsumexp, do, di, - bq=bq_dq, - bkv=bkv_dq, + bq=bq_dq, # pyrefly: ignore[bad-argument-type] + bkv=bkv_dq, # pyrefly: ignore[bad-argument-type] is_mqa=is_mqa, - mask_info=dq_mask_info, + mask_info=dq_mask_info, # pyrefly: ignore[bad-argument-type] mask_value=mask_value, attn_logits_soft_cap=attn_logits_soft_cap, q_layout=block_sizes.q_layout, @@ -2348,7 +2348,7 @@ def _splash_attention_manual_bwd( save_residuals=save_residuals, mask_value=mask_value, is_mqa=is_mqa, - block_sizes=block_sizes, + block_sizes=block_sizes, # pyrefly: ignore[bad-argument-type] residual_checkpoint_name=residual_checkpoint_name, mask_function=mask_function, attn_logits_soft_cap=attn_logits_soft_cap, @@ -2422,11 +2422,11 @@ def manual_sharding_spec(self, sharding: jax.sharding.NamedSharding): # Shard q_sequence over the sequence dimension only. q_sequence_spec = jax.sharding.PartitionSpec(spec[1]) mask_info_specs = mask_info_lib.MaskInfo( # pytype: disable=wrong-arg-types - data_next=spec if self.fwd_mask_info.data_next is not None else None, - mask_next=spec if self.fwd_mask_info.mask_next is not None else None, - block_mask=spec if self.fwd_mask_info.block_mask is not None else None, - partial_mask_blocks=partial_mask_blocks_spec if self.fwd_mask_info.partial_mask_blocks is not None else None, - q_sequence=q_sequence_spec if self.fwd_mask_info.q_sequence is not None else None, + data_next=spec if self.fwd_mask_info.data_next is not None else None, # pyrefly: ignore[bad-argument-type] + mask_next=spec if self.fwd_mask_info.mask_next is not None else None, # pyrefly: ignore[bad-argument-type] + block_mask=spec if self.fwd_mask_info.block_mask is not None else None, # pyrefly: ignore[bad-argument-type] + partial_mask_blocks=partial_mask_blocks_spec if self.fwd_mask_info.partial_mask_blocks is not None else None, # pyrefly: ignore[bad-argument-type] + q_sequence=q_sequence_spec if self.fwd_mask_info.q_sequence is not None else None, # pyrefly: ignore[bad-argument-type] ) return SplashAttentionKernel( mask_info_specs, diff --git a/src/maxtext/kernels/gather_reduce_sc.py b/src/maxtext/kernels/gather_reduce_sc.py index 5b3b8e7597..792f799e33 100644 --- a/src/maxtext/kernels/gather_reduce_sc.py +++ b/src/maxtext/kernels/gather_reduce_sc.py @@ -473,7 +473,7 @@ def perform_add( is_odd = arith.cmpi( arith.CmpIPredicate.eq, - parity, + parity, # pyrefly: ignore[bad-argument-type] const_lut(0), ) raw_weights_f32 = arith.select(is_odd, weights_odds, weights_evens) @@ -533,8 +533,8 @@ def get_row_val(row_idx): _BF16[2, 16], scratch_local, [ - arith.addi(const_lut(row_idx_l), row_add), - arith.muli(col_pos, const_lut(2)), # NOTYPO + arith.addi(const_lut(row_idx_l), row_add), # pyrefly: ignore[bad-argument-type] + arith.muli(col_pos, const_lut(2)), # NOTYPO # pyrefly: ignore[bad-argument-type] ], enable_all_sublanes_mask, ) @@ -567,20 +567,20 @@ def get_row_val(row_idx): row0 = get_row_val(0) if weights_local is not None: - row0 = arith.mulf(row0, weights_vecs[0]) + row0 = arith.mulf(row0, weights_vecs[0]) # pyrefly: ignore[unsupported-operation] if topk_wgt_zero_nan: row0 = arith.select( - arith.cmpf(arith.CmpFPredicate.OEQ, weights_vecs[0], zero_vec_f32), + arith.cmpf(arith.CmpFPredicate.OEQ, weights_vecs[0], zero_vec_f32), # pyrefly: ignore[unbound-name, unsupported-operation] zero_vec_f32, row0, ) row8 = get_row_val(8) if weights_local is not None: - row8 = arith.mulf(row8, weights_vecs[8]) + row8 = arith.mulf(row8, weights_vecs[8]) # pyrefly: ignore[unsupported-operation] if topk_wgt_zero_nan: row8 = arith.select( - arith.cmpf(arith.CmpFPredicate.OEQ, weights_vecs[8], zero_vec_f32), + arith.cmpf(arith.CmpFPredicate.OEQ, weights_vecs[8], zero_vec_f32), # pyrefly: ignore[unsupported-operation] zero_vec_f32, row8, ) @@ -588,12 +588,12 @@ def get_row_val(row_idx): for sum_idx in range(7): tmp_row0 = get_row_val(sum_idx + 1) if weights_local is not None: - tmp_row0 = arith.mulf(tmp_row0, weights_vecs[sum_idx + 1]) + tmp_row0 = arith.mulf(tmp_row0, weights_vecs[sum_idx + 1]) # pyrefly: ignore[unsupported-operation] if topk_wgt_zero_nan: tmp_row0 = arith.select( arith.cmpf( arith.CmpFPredicate.OEQ, - weights_vecs[sum_idx + 1], + weights_vecs[sum_idx + 1], # pyrefly: ignore[unsupported-operation] zero_vec_f32, ), zero_vec_f32, @@ -604,12 +604,12 @@ def get_row_val(row_idx): tmp_row8 = get_row_val(8 + sum_idx + 1) if weights_local is not None: - tmp_row8 = arith.mulf(tmp_row8, weights_vecs[8 + sum_idx + 1]) + tmp_row8 = arith.mulf(tmp_row8, weights_vecs[8 + sum_idx + 1]) # pyrefly: ignore[unsupported-operation] if topk_wgt_zero_nan: tmp_row8 = arith.select( arith.cmpf( arith.CmpFPredicate.OEQ, - weights_vecs[8 + sum_idx + 1], + weights_vecs[8 + sum_idx + 1], # pyrefly: ignore[unsupported-operation] zero_vec_f32, ), zero_vec_f32, @@ -743,7 +743,7 @@ def fill_out_offset_tile(offset_tile_out_local, col_pos, row_pos=None): ) if row_pos is not None: - local_vec = arith.addi(vec, row_vec_offset) + local_vec = arith.addi(vec, row_vec_offset) # pyrefly: ignore[unbound-name] else: local_vec = vec @@ -1154,7 +1154,7 @@ def fill_out_offset_tile(offset_tile_out_local, col_pos, row_pos=None): lower_bound=const_lut(row_chunk_size * 1), upper_bound=const_lut(row_chunk_size * ((idx.shape[0] // num_sc // row_chunk_size) - 1)), step=const_lut(row_chunk_size * 2), - iter_args=[ + iter_args=[ # pyrefly: ignore[bad-argument-type] scratch_0, scratch_out_0, sflag_0, diff --git a/src/maxtext/kernels/megablox/backend.py b/src/maxtext/kernels/megablox/backend.py index e736bdb000..de3fab8b92 100644 --- a/src/maxtext/kernels/megablox/backend.py +++ b/src/maxtext/kernels/megablox/backend.py @@ -345,7 +345,7 @@ def gmm( group_offset = group_offset[None] num_current_groups = rhs.shape[0] num_total_groups = group_sizes.shape[0] - group_sizes = _validate_args(lhs=lhs, rhs=rhs, group_sizes=group_sizes) + group_sizes = _validate_args(lhs=lhs, rhs=rhs, group_sizes=group_sizes) # pyrefly: ignore[bad-argument-type] # Gather shape information. m, k, n = (lhs.shape[0], lhs.shape[1], rhs.shape[2]) @@ -607,7 +607,7 @@ def tgmm( group_offset = jnp.array([0], dtype=jnp.int32) else: group_offset = group_offset[None] - group_sizes = _validate_args(lhs=lhs, rhs=rhs, group_sizes=group_sizes, expected_rhs_dims=2) + group_sizes = _validate_args(lhs=lhs, rhs=rhs, group_sizes=group_sizes, expected_rhs_dims=2) # pyrefly: ignore[bad-argument-type] # Gather shape information. k, m, n = (lhs.shape[0], lhs.shape[1], rhs.shape[1]) diff --git a/src/maxtext/kernels/megablox/ops.py b/src/maxtext/kernels/megablox/ops.py index b7487c7368..06f9aa9eac 100644 --- a/src/maxtext/kernels/megablox/ops.py +++ b/src/maxtext/kernels/megablox/ops.py @@ -63,8 +63,8 @@ def gmm( existing_out: jnp.ndarray | None = None, transpose_rhs: bool = False, interpret: bool = False, - lhs_quantize_dtype: Literal[jnp.int4, jnp.int8] | None = None, - rhs_quantize_dtype: Literal[jnp.int4, jnp.int8] | None = None, + lhs_quantize_dtype: Literal[jnp.int4, jnp.int8] | None = None, # pyrefly: ignore[invalid-literal] + rhs_quantize_dtype: Literal[jnp.int4, jnp.int8] | None = None, # pyrefly: ignore[invalid-literal] use_qwix_quantization: bool = False, use_tokamax_backend: bool = False, weight_gather_axes: List[Tuple[str, int]] | None = None, @@ -154,15 +154,15 @@ def _gmm_fwd( """Forward function for GMM VJP.""" if quantization_rule: if quantization_rule.act_qtype and not isinstance(lhs, qpl.QArray): - lhs = qpl.quantize( + lhs = qpl.quantize( # pyrefly: ignore[bad-assignment] lhs, quantization_rule.act_qtype, channelwise_axes=[] if quantization_rule.disable_channelwise_axes else [0], - calibration_method=quantization_rule.act_calibration_method, + calibration_method=quantization_rule.act_calibration_method, # pyrefly: ignore[bad-argument-type] ) if quantization_rule.weight_qtype and not isinstance(rhs, qpl.QArray): if not use_manual_quantization: - rhs = qpl.quantize( + rhs = qpl.quantize( # pyrefly: ignore[bad-assignment] rhs, quantization_rule.weight_qtype, # If only considering the fwd pass, we could also enable channelwise @@ -172,7 +172,7 @@ def _gmm_fwd( calibration_method=quantization_rule.weight_calibration_method, ) else: - rhs = quantizations.manual_quantize( + rhs = quantizations.manual_quantize( # pyrefly: ignore[bad-assignment] rhs, quantization_rule.weight_qtype, calibration_method=quantization_rule.weight_calibration_method, diff --git a/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py b/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py index 8240e1bf14..c1ebeda0e3 100644 --- a/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py +++ b/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_gmm_kernel.py @@ -210,7 +210,7 @@ class GmmConfigs: @property def num_quant_blocks_per_tile_k(self) -> int: - return pl.cdiv(self.tiles.tile_k, self.rhs_cfgs.quant_block_size) + return pl.cdiv(self.tiles.tile_k, self.rhs_cfgs.quant_block_size) # pyrefly: ignore[no-matching-overload] @property def out_size_n(self) -> int: @@ -253,7 +253,7 @@ def rhs_scale_index_map(self, n_id: jax.Array, gm_id: jax.Array, k_id: jax.Array # Simply multiplying k_id by num_quant_blocks_per_tile_k will not work # since a single quant block could be shared along multiple k tile. k_row = k_id * self.cfgs.tiles.tile_k - b_row = k_row // self.cfgs.rhs_cfgs.quant_block_size + b_row = k_row // self.cfgs.rhs_cfgs.quant_block_size # pyrefly: ignore[unsupported-operation] b_tile_id = b_row // self.cfgs.num_quant_blocks_per_tile_k return (group_id, b_tile_id, 0, n_id) @@ -420,8 +420,8 @@ def _matmul(is_first_k_step: bool, is_last_k_step: bool): acc_n = jnp.zeros((cfgs.tiles.tile_m, col_size), dtype=acc_ref.dtype) for b_id in range(cfgs.num_quant_blocks_per_tile_k): - start_k = b_id * rhs_qbs - end_k = start_k + rhs_qbs + start_k = b_id * rhs_qbs # pyrefly: ignore[unsupported-operation] + end_k = start_k + rhs_qbs # pyrefly: ignore[unsupported-operation] block_acc = jnp.matmul( tiled_lhs[:, start_k:end_k], @@ -458,8 +458,8 @@ def _matmul(is_first_k_step: bool, is_last_k_step: bool): col_size = end_n - start_n acc_n = jnp.zeros((cfgs.tiles.tile_m, col_size), dtype=acc_ref.dtype) - for start_k in range(0, cfgs.tiles.tile_k, q_block_size): - end_k = min(cfgs.tiles.tile_k, start_k + q_block_size) + for start_k in range(0, cfgs.tiles.tile_k, q_block_size): # pyrefly: ignore[bad-argument-type] + end_k = min(cfgs.tiles.tile_k, start_k + q_block_size) # pyrefly: ignore[unsupported-operation] block_lhs = tiled_lhs[:, start_k:end_k] block_rhs = tiled_rhs[start_k:end_k, start_n:end_n] @@ -493,7 +493,7 @@ def _matmul(is_first_k_step: bool, is_last_k_step: bool): # Apply rhs subchannel scale per quant block. if cfgs.rhs_cfgs.should_dequantize_after_matmul: - b_id = start_k // cfgs.rhs_cfgs.quant_block_size + b_id = start_k // cfgs.rhs_cfgs.quant_block_size # pyrefly: ignore[unsupported-operation] rhs_scale_slice = tiled_rhs_ref.get_scale() block_acc *= rhs_scale_slice[b_id, :, start_n:end_n].astype(acc_ref.dtype) @@ -510,7 +510,7 @@ def _matmul(is_first_k_step: bool, is_last_k_step: bool): tiled_rhs_bias = tiled_rhs_ref.get_bias() acc += tiled_rhs_bias.astype(acc.dtype) if cfgs.has_partial_sum: - ps_tile = tiled_ps_ref[...].reshape(acc.shape) + ps_tile = tiled_ps_ref[...].reshape(acc.shape) # pyrefly: ignore[unsupported-operation] acc += ps_tile.astype(acc.dtype) acc = apply_act_fn(acc, cfgs.fuse_act) @@ -829,8 +829,8 @@ def kernel_main( if cfgs.zero_init: zero_size = zero_out_start( out_ref, - zero_ref, - semaphore_ref, + zero_ref, # pyrefly: ignore[bad-argument-type] + semaphore_ref, # pyrefly: ignore[bad-argument-type] metadata_ref, num_gm, dims=cfgs.dims, @@ -840,7 +840,7 @@ def kernel_main( if cfgs.fuse_act is not None: rhs_up_ref = jax.tree.map(lambda x: x.at[..., cfgs.out_size_n :], rhs_ref) - rhs_ref = FusedWeightsRef(gate=rhs_ref, up=rhs_up_ref) + rhs_ref = FusedWeightsRef(gate=rhs_ref, up=rhs_up_ref) # pyrefly: ignore[bad-assignment] rhs_spec = FusedWeightsRef( gate=rhs_spec, @@ -866,7 +866,7 @@ def kernel_main( pipeline_fn(lhs_in, rhs_ref, ps_in, out_in, scratches=scratches) if cfgs.zero_init: - zero_out_end(out_ref, semaphore_ref, zero_size, dims=cfgs.dims) + zero_out_end(out_ref, semaphore_ref, zero_size, dims=cfgs.dims) # pyrefly: ignore[bad-argument-type, unbound-name] def calculate_tiling( @@ -909,7 +909,7 @@ def calculate_tiling( tile_n_limit //= fuse_act_factor def _is_tile_k_quant_block_compatible(tk: int) -> bool: - if tk % rhs_cfgs.quant_block_size != 0 and rhs_cfgs.quant_block_size % tk != 0: + if tk % rhs_cfgs.quant_block_size != 0 and rhs_cfgs.quant_block_size % tk != 0: # pyrefly: ignore[unsupported-operation] return False return True @@ -1046,7 +1046,7 @@ def get_cost_estimate(cfgs: GmmConfigs): rhs_size = dims.size_group * dims.size_k * dims.size_n rhs_bytes = rhs_size * rhs_bits // 8 if cfgs.rhs_cfgs.has_scale: - num_quant_blocks = pl.cdiv(dims.size_k, cfgs.rhs_cfgs.quant_block_size) + num_quant_blocks = pl.cdiv(dims.size_k, cfgs.rhs_cfgs.quant_block_size) # pyrefly: ignore[no-matching-overload] rhs_bytes += dims.size_group * num_quant_blocks * dims.size_n * fp32_bytes if cfgs.rhs_cfgs.has_bias: rhs_bytes += dims.size_group * dims.size_n * fp32_bytes @@ -1113,14 +1113,14 @@ def make_gmm_configs( lhs_q_dtype = None if maybe_quantize_lhs and rhs_cfgs.should_dequantize_after_matmul: # Choose lhs quantization dtype based on TPU hardware support. - is_rhs_float = jnp.issubdtype(rhs_quant_dtype, jnp.floating) + is_rhs_float = jnp.issubdtype(rhs_quant_dtype, jnp.floating) # pyrefly: ignore[bad-argument-type] tpu_info = pltpu.get_tpu_info() # Check if there is hardware compute support for rhs dtype group. if tpu_info.fp8_ops_per_second > 0: # Special handling for 4-bit integer rhs as it can be converted to fp8 # without a numeric issues. Note that this is not the case for 4-bit # floating rhs as conversion to int8 will cause numeric issues. - is_rhs_4bits = jax.dtypes.itemsize_bits(rhs_quant_dtype) == 4 + is_rhs_4bits = jax.dtypes.itemsize_bits(rhs_quant_dtype) == 4 # pyrefly: ignore[bad-argument-type] if is_rhs_float or is_rhs_4bits: lhs_q_dtype = jnp.float8_e4m3fn.dtype if tpu_info.int8_ops_per_second > 0: @@ -1199,7 +1199,7 @@ def gmm_v2( partial_sum: jax.Array | None = None, # [size_m, size_n] group_offset: jax.Array | None = None, # int32[1] *, - tile_info: TileSizes | TileFn = calculate_tiling, + tile_info: TileSizes | TileFn = calculate_tiling, # pyrefly: ignore[bad-function-definition] vmem_limit_bytes: int | None = None, precision: jax.lax.Precision = jax.lax.Precision.DEFAULT, preferred_element_type: jnp.dtype | None = None, @@ -1348,7 +1348,7 @@ def gmm_v2( num_scalar_prefetch=2, in_specs=in_specs, out_specs=pl.BlockSpec(memory_space=pltpu.HBM), - scratch_shapes=scratch_shapes, + scratch_shapes=scratch_shapes, # pyrefly: ignore[bad-argument-type] ), compiler_params=pltpu.CompilerParams( vmem_limit_bytes=vmem_limit_bytes, diff --git a/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_tgmm_kernel.py b/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_tgmm_kernel.py index 01b17bd247..6b653b5ae7 100644 --- a/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_tgmm_kernel.py +++ b/src/maxtext/kernels/megablox/pallas_mosaic_tpu_v2_tgmm_kernel.py @@ -54,6 +54,7 @@ class OperandRef: jnp.dtype, jnp.dtype, int, + bool, ], gmm_v2.TileSizes, ] @@ -89,6 +90,7 @@ def calculate_tgmm_tiling( out_dtype: jnp.dtype, acc_dtype: jnp.dtype, target_zero_ref_bytes: int, + has_partial_sum: bool = False, ) -> gmm_v2.TileSizes: """Calculate optimal tile sizes for TGMM kernel.""" # In tgmm, we calculate lhs.T @ dout which doesn't require quantization. @@ -121,8 +123,10 @@ def within_vmem_limit(tile_m, tile_k, tile_n): # XLU's transpose. in order to reduce redundant XLU computation, instead # of performing XLU's transpose every time lhs is pushed into XLU, it # caches the transposed value into VMEM. this increases VMEM requirement. + ps_bytes = tile_k * tile_n * num_buffers * out_bytes if has_partial_sum else 0 budget = ( tile_k * tile_n * (acc_bytes + num_buffers * out_bytes) + + ps_bytes + (num_buffers + 1) * (tile_m * tile_k * lhs_bytes) + num_buffers * (tile_m * tile_n * rhs_bytes) # Reserve VMEM for zero_ref. Use the upper bound target_zero_ref_bytes @@ -179,6 +183,7 @@ def make_tgmm_configs( lhs: jax.Array, # [m, k] rhs: jax.Array, # [m, n] rhs_scale: jax.Array, # [1, 1, n] (per-N scale) + partial_sum: jax.Array | None, group_sizes: jax.Array, num_actual_groups: int, *, @@ -252,10 +257,11 @@ def make_tgmm_configs( dims, lhs_cfgs, rhs_cfgs, - vmem_limit_bytes, + vmem_limit_bytes, # pyrefly: ignore[bad-argument-type] out_dtype, acc_dtype, target_zero_ref_bytes, + partial_sum is not None, ) return gmm_v2.GmmConfigs( @@ -263,7 +269,7 @@ def make_tgmm_configs( tiles=tiles, lhs_cfgs=lhs_cfgs, rhs_cfgs=rhs_cfgs, - has_partial_sum=False, # This should always be False until partial sum support is added in bwd pass. + has_partial_sum=(partial_sum is not None), out_dtype=jnp.dtype(out_dtype), acc_dtype=jnp.dtype(acc_dtype), # GMM's 'zero_init' zeros unvisited m-rows via DMA, which doesn't apply to @@ -280,6 +286,7 @@ def tgmm_inner_kernel( tiled_rhs_ref: OperandRef, # .value: [tile_m // size_lhs_sublane, size_lhs_sublane, tile_n] # .scale: [1, 1, tile_n] or None + tiled_ps_ref: jax.Array | None, tiled_out_ref: jax.Array, acc_ref: jax.Array, metadata_ref: gmm_v2.MetadataRef, @@ -340,8 +347,10 @@ def _matmul(is_new_group: bool, is_group_changing: bool): if is_group_changing: if cfgs.rhs_cfgs.has_scale: - scale_slice = tiled_rhs_scale_ref[0] + scale_slice = tiled_rhs_scale_ref[0] # pyrefly: ignore[unsupported-operation] acc *= scale_slice + if cfgs.has_partial_sum: + acc += tiled_ps_ref[...].astype(acc.dtype) tiled_out_ref[...] = acc.astype(tiled_out_ref.dtype) else: acc_ref[...] = acc @@ -430,7 +439,7 @@ def out_index_map(self, n_id: jax.Array, k_id: jax.Array, gm_id: jax.Array): def generate_tgmm_block_specs( metadata_ref: gmm_v2.MetadataRef, cfgs: gmm_v2.GmmConfigs -) -> Tuple[Tuple[pl.BlockSpec, OperandRef], pl.BlockSpec]: +) -> Tuple[Tuple[pl.BlockSpec, OperandRef, pl.BlockSpec | None], pl.BlockSpec]: """Generates block specs for the given lhs, rhs, and out refs.""" index_map = TgmmIndexMaps(metadata_ref, cfgs) # NB: in tgmm, LHS is reshaped from (M, K) to (-1, size_lhs_sublane, K) so @@ -457,8 +466,14 @@ def generate_tgmm_block_specs( (None, cfgs.tiles.tile_k, cfgs.tiles.tile_n), index_map.out_index_map, ) - - return (lhs_block_spec, rhs_spec), out_block_spec + ps_block_spec = None + if cfgs.has_partial_sum: + ps_block_spec = pl.BlockSpec( + (None, cfgs.tiles.tile_k, cfgs.tiles.tile_n), + index_map.out_index_map, + ) + in_specs = (lhs_block_spec, rhs_spec, ps_block_spec) + return in_specs, out_block_spec def zero_out_start( @@ -526,6 +541,7 @@ def tgmm_kernel_main( group_offset_ref, # int32[1] lhs_ref, # [m, k] rhs_ref, # OperandRef: .value [m, n], .scale [1, 1, n] or None + partial_sum_ref, # [num_actual_groups, k, n] or None out_ref, # [num_actual_groups, k, n] # Scratch memory acc_ref: jax.Array, # [tile_k, tile_n] @@ -578,9 +594,12 @@ def tgmm_kernel_main( rhs_value = rhs_ref.value rhs_in = rhs_value.reshape(-1, cfgs.dims.size_lhs_sublane, rhs_value.shape[-1]) rhs_operand = OperandRef(value=rhs_in, scale=rhs_ref.scale) + ps_in = None + if cfgs.has_partial_sum: + ps_in = partial_sum_ref scratches = [acc_ref, metadata_ref] - pipeline_fn(lhs_in, rhs_operand, out_ref, scratches=scratches) + pipeline_fn(lhs_in, rhs_operand, ps_in, out_ref, scratches=scratches) zero_out_end( num_groups_to_zero, out_ref, @@ -632,6 +651,7 @@ def tgmm_v2( group_sizes: jax.Array, num_actual_groups: int, rhs_scale: jax.Array | None = None, # [1, 1, size_n] (per-N scale) + partial_sum: jax.Array | None = None, group_offset: jax.Array | None = None, *, tile_info: gmm_v2.TileSizes | TileTgmmFn = calculate_tgmm_tiling, @@ -682,12 +702,13 @@ def tgmm_v2( cfgs = make_tgmm_configs( lhs, rhs, - rhs_scale, + rhs_scale, # pyrefly: ignore[bad-argument-type] + partial_sum, group_sizes, num_actual_groups, tile_info=tile_info, vmem_limit_bytes=vmem_limit_bytes, - out_dtype=preferred_element_type, + out_dtype=preferred_element_type, # pyrefly: ignore[bad-argument-type] acc_dtype=acc_dtype, target_zero_ref_bytes=target_zero_ref_bytes, ) @@ -729,23 +750,27 @@ def tgmm_v2( pad_n = aligned_n - dims.size_n if pad_n > 0: rhs_scale = jnp.pad(rhs_scale, ((0, 0), (0, 0), (0, pad_n))) - rhs = OperandRef(value=rhs, scale=rhs_scale) + rhs = OperandRef(value=rhs, scale=rhs_scale) # pyrefly: ignore[bad-assignment] hbm_spec = pl.BlockSpec(memory_space=pltpu.HBM) + partial_sum_spec = None + if partial_sum is not None: + partial_sum_spec = hbm_spec in_specs = [ hbm_spec, # lhs # the tree.map build a # OperandRef(value=hbm_spec, scale=None if scale is None else hbm_spec. jax.tree.map(lambda _: hbm_spec, rhs), # rhs + partial_sum_spec, ] - return pl.pallas_call( + raw_out = pl.pallas_call( functools.partial(tgmm_kernel_main, cfgs=cfgs), out_shape=out_init, grid_spec=pltpu.PrefetchScalarGridSpec( num_scalar_prefetch=2, in_specs=in_specs, out_specs=pl.BlockSpec(memory_space=pltpu.HBM), - scratch_shapes=scratch_shapes, + scratch_shapes=scratch_shapes, # pyrefly: ignore[bad-argument-type] ), compiler_params=pltpu.CompilerParams( vmem_limit_bytes=vmem_limit_bytes, @@ -756,4 +781,10 @@ def tgmm_v2( # the metadata here is for profiling, debugging, and cost modeling. # It does not affect the kernel's computation. metadata=gmm_v2.get_metadata(cfgs), - )(group_sizes, group_offset, lhs, rhs)[:, : dims.size_k, : dims.size_n] + )(group_sizes, group_offset, lhs, rhs, partial_sum)[:, : dims.size_k, : dims.size_n] + + if partial_sum is not None: + local_group_sizes = lax.dynamic_slice(group_sizes, (group_offset[0],), (num_actual_groups,)) + empty_mask = (local_group_sizes == 0).reshape(num_actual_groups, 1, 1) + return jnp.where(empty_mask, partial_sum, raw_out) + return raw_out diff --git a/src/maxtext/kernels/ragged/ragged_gather.py b/src/maxtext/kernels/ragged/ragged_gather.py index 14aa989c88..fc9aba84b8 100644 --- a/src/maxtext/kernels/ragged/ragged_gather.py +++ b/src/maxtext/kernels/ragged/ragged_gather.py @@ -140,8 +140,8 @@ def _(): # but for other bitwidths, this is not possible. Therefore, we bitcast data # into 32-bits first to fetch packing number of rows per dma and later # perform bitwise unpacking / packing to obtain desired results. - in_32b_hbm_ref = in_hbm_ref.bitcast(jnp.uint32) - out_32b_hbm_ref = out_hbm_ref.bitcast(jnp.uint32) + in_32b_hbm_ref = in_hbm_ref.bitcast(jnp.uint32) # pyrefly: ignore[missing-attribute] + out_32b_hbm_ref = out_hbm_ref.bitcast(jnp.uint32) # pyrefly: ignore[missing-attribute] for col_vmem_start in range(0, col_size, num_lanes): col_hbm_start = col_tile_start + col_vmem_start @@ -333,7 +333,7 @@ def _fallback_implementation( """Fallback to (non-ragged) JAX implementation for ragged gather.""" out = x[indices] if has_weights: - out = out * weights[:, None] + out = out * weights[:, None] # pyrefly: ignore[unsupported-operation] return out @@ -433,7 +433,7 @@ def ragged_gather( indices = jnp.pad(indices, ((0, out_pad_size))) if has_weights: - weights = jnp.pad(weights, ((0, out_pad_size)), constant_values=1.0) + weights = jnp.pad(weights, ((0, out_pad_size)), constant_values=1.0) # pyrefly: ignore[bad-argument-type] else: # Provide a dummy weights array; the kernel won't use it. weights = jnp.ones((out_size + out_pad_size,), dtype=jnp.float32) @@ -454,7 +454,7 @@ def ragged_gather( has_weights=has_weights, ), compiler_params=pltpu.CompilerParams( # pytype: disable=wrong-keyword-args - **_COMPILER_PARAMS, + **_COMPILER_PARAMS, # pyrefly: ignore[bad-argument-type] ), cost_estimate=get_cost_estimate( out_size=out_size + out_pad_size, @@ -466,7 +466,7 @@ def ragged_gather( ), mesh=vector_mesh, name="sc_ragged_gather", - **{ + **{ # pyrefly: ignore[bad-argument-type] _OUT_KW: jax.ShapeDtypeStruct((out_size + out_pad_size, aligned_hidden_size), dtype), _SCRATCH_KW: [ pltpu.VMEM((num_simd_lanes,), jnp.int32), diff --git a/src/maxtext/kernels/ragged/ragged_gather_reduce.py b/src/maxtext/kernels/ragged/ragged_gather_reduce.py index d5bb3314b6..d9385ce9d7 100644 --- a/src/maxtext/kernels/ragged/ragged_gather_reduce.py +++ b/src/maxtext/kernels/ragged/ragged_gather_reduce.py @@ -228,8 +228,8 @@ def row_loop(row_block_id): input_dtype_bits = jax.dtypes.itemsize_bits(in_dtype) input_packing = 32 // input_dtype_bits - in_32b_hbm_ref = in_hbm_ref.bitcast(jnp.uint32) - out_32b_hbm_ref = out_hbm_ref.bitcast(jnp.uint32) + in_32b_hbm_ref = in_hbm_ref.bitcast(jnp.uint32) # pyrefly: ignore[missing-attribute] + out_32b_hbm_ref = out_hbm_ref.bitcast(jnp.uint32) # pyrefly: ignore[missing-attribute] for col_vmem_start in range(0, col_size, num_lanes): col_hbm_start = col_start + col_vmem_start @@ -559,7 +559,7 @@ def ragged_gather_reduce( num_column_partitions=num_column_partitions, ), compiler_params=pltpu.CompilerParams( # pytype: disable=wrong-keyword-args - **_COMPILER_PARAMS, + **_COMPILER_PARAMS, # pyrefly: ignore[bad-argument-type] ), cost_estimate=get_cost_estimate( padded_input_size=padded_input_size, @@ -571,7 +571,7 @@ def ragged_gather_reduce( ), mesh=vector_mesh, name="sc_ragged_gather_reduce", - **{ + **{ # pyrefly: ignore[bad-argument-type] _OUT_KW: jax.ShapeDtypeStruct( (padded_input_size // reduce_group_size, aligned_hidden_size), jnp.float32, diff --git a/src/maxtext/layers/attention_op.py b/src/maxtext/layers/attention_op.py index 09c467706a..1c015b464b 100644 --- a/src/maxtext/layers/attention_op.py +++ b/src/maxtext/layers/attention_op.py @@ -615,6 +615,7 @@ def generate_attention_mask( previous_chunk: Any = None, bidirectional_mask: Any = None, compressed_mask: Optional[Array] = None, + segment_positions: Array | None = None, ) -> Array | None: """Generates a combined attention mask for Transformer models. @@ -675,6 +676,9 @@ def generate_attention_mask( compressed_mask: Optional `Array`. A pre-computed attention mask for compressed kv blocks (e.g., DeepSeek-V4 compressed attention). If provided, it is concatenated with the dynamically generated uncompressed mask. + segment_positions: Optional `Array` of shape `[batch_size, + q_sequence_length]`. Identifies original positions for load-balanced + context-parallel inputs. Returns: An `Array` representing the attention mask, with shape @@ -710,20 +714,33 @@ def generate_attention_mask( elif model_mode == MODEL_MODE_AUTOREGRESSIVE and q_seq_len == 1: # In autoregression, the query position is the last position in the KV sequence. next_pos = kv_seq_len - 1 + use_segment_positions = ( + segment_positions is not None + and self.config.context_parallel_load_balance + and self.mesh.shape.get(self.config.context_sharding, 1) > 1 + and previous_chunk is None + and model_mode != MODEL_MODE_AUTOREGRESSIVE + ) + if use_segment_positions: + position_row_ids = segment_positions[:, :, None] + position_col_ids = segment_positions[:, None, :] causal_mask = None if model_mode != MODEL_MODE_AUTOREGRESSIVE and self.attention_type not in ( AttentionType.FULL, AttentionType.COMPRESSED, ): - mask_shape = (q_seq_len, kv_seq_len) - # row_ids indicates the position of query - # col_ids indicates the position of kv - row_ids = jax.lax.broadcasted_iota(jnp.int32, mask_shape, 0) - col_ids = jax.lax.broadcasted_iota(jnp.int32, mask_shape, 1) - # Attention mask for chunked prefill is generated in the same way - # as mentioned in SARATHI - https://arxiv.org/abs/2308.16369 - causal_mask = (col_ids <= row_ids + next_pos)[None, None, None, :, :] + if use_segment_positions: + causal_mask = (position_col_ids <= position_row_ids)[:, None, None, :, :] + else: + mask_shape = (q_seq_len, kv_seq_len) + # row_ids indicates the position of query + # col_ids indicates the position of kv + row_ids = jax.lax.broadcasted_iota(jnp.int32, mask_shape, 0) + col_ids = jax.lax.broadcasted_iota(jnp.int32, mask_shape, 1) + # Attention mask for chunked prefill is generated in the same way + # as mentioned in SARATHI - https://arxiv.org/abs/2308.16369 + causal_mask = (col_ids <= row_ids + next_pos)[None, None, None, :, :] output_mask = None if (mask is not None) and (causal_mask is not None): @@ -737,11 +754,17 @@ def generate_attention_mask( if self.sliding_window_size is None: raise ValueError("Sliding_window_size must be set if Local Sliding attention type") - row_ids_sliding = jax.lax.broadcasted_iota(jnp.int32, (q_seq_len, 1), 0) + next_pos - col_ids_sliding = jax.lax.broadcasted_iota(jnp.int32, (1, kv_seq_len), 1) + if use_segment_positions: + row_ids_sliding = position_row_ids + col_ids_sliding = position_col_ids + else: + row_ids_sliding = jax.lax.broadcasted_iota(jnp.int32, (q_seq_len, 1), 0) + next_pos + col_ids_sliding = jax.lax.broadcasted_iota(jnp.int32, (1, kv_seq_len), 1) sliding_mask = (col_ids_sliding > (row_ids_sliding - self.sliding_window_size)) & ( col_ids_sliding <= row_ids_sliding ) + if use_segment_positions: + sliding_mask = sliding_mask[:, None, None, :, :] output_mask = sliding_mask * output_mask elif self.attention_type == AttentionType.COMPRESSED: if compressed_mask is None: @@ -772,12 +795,17 @@ def generate_attention_mask( return jnp.concatenate([uncompressed_mask, compressed_mask], axis=-1) elif self.attention_type == AttentionType.CHUNK and output_mask is not None: - mask_shape = (q_seq_len, kv_seq_len) - chunk_mask = _generate_chunk_attention_mask( - mask_shape=(q_seq_len, kv_seq_len), - chunk_size=self.chunk_attn_window_size, - q_offset=next_pos, - ) + if use_segment_positions: + same_chunk = (position_row_ids // self.chunk_attn_window_size) == ( + position_col_ids // self.chunk_attn_window_size + ) + chunk_mask = (same_chunk & (position_row_ids >= position_col_ids))[:, None, None, :, :] + else: + chunk_mask = _generate_chunk_attention_mask( + mask_shape=(q_seq_len, kv_seq_len), + chunk_size=self.chunk_attn_window_size, + q_offset=next_pos, + ) output_mask = chunk_mask * output_mask if bidirectional_mask is not None: @@ -976,6 +1004,7 @@ def apply_attention( decoder_segment_ids, model_mode, previous_chunk, + segment_positions=segment_positions, bidirectional_mask=bidirectional_mask, sinks=sinks, indexer_mask=indexer_mask, @@ -1275,28 +1304,43 @@ def create_sa_config(config, query, key, attn_logits_soft_cap): else: mask = mask_module.CausalMask(shape=mask_shape) - # Create LoadBalancedCausalMask if cp and load_balancing - if cp_size > 1 and load_balanced_context_parallel: + use_load_balanced_cp = cp_size > 1 and load_balanced_context_parallel + if use_load_balanced_cp and self.attention_type != AttentionType.FULL: mask = LoadBalancedCausalMask(shape=mask_shape, cp_size=cp_size) - # TODO: figure out local_sliding attention + load_balancing, default is global # Apply local masking if local sliding attention is enabled. if self.attention_type == AttentionType.LOCAL_SLIDING: if self.sliding_window_size is None: raise ValueError("Sliding_window_size must be set if Local Sliding attention type") - mask &= mask_module.LocalMask( - shape=(query.shape[2], key.shape[2]), - window_size=(self.sliding_window_size - 1, self.sliding_window_size), - offset=0, - ) + local_window_size = (self.sliding_window_size - 1, self.sliding_window_size) + if use_load_balanced_cp: + mask &= LoadBalancedLocalMask( + shape=(query.shape[2], key.shape[2]), + window_size=local_window_size, + offset=0, + cp_size=cp_size, + ) + else: + mask &= mask_module.LocalMask( + shape=(query.shape[2], key.shape[2]), + window_size=local_window_size, + offset=0, + ) elif self.attention_type == AttentionType.CHUNK: if self.chunk_attn_window_size is None: raise ValueError("chunk_attn_window_size must be set for chunk attention type") - mask &= ChunkedCausalMask( - shape=(query.shape[2], key.shape[2]), - chunk_size=self.chunk_attn_window_size, - ) + if use_load_balanced_cp: + mask &= LoadBalancedChunkedCausalMask( + shape=(query.shape[2], key.shape[2]), + chunk_size=self.chunk_attn_window_size, + cp_size=cp_size, + ) + else: + mask &= ChunkedCausalMask( + shape=(query.shape[2], key.shape[2]), + chunk_size=self.chunk_attn_window_size, + ) max_logit_value = None if self.config.use_tokamax_splash: @@ -1500,6 +1544,7 @@ def kernel_fn(q, k, v, d, s): block_q=self.block_q, mask=materialized_mask, mask_value=DEFAULT_MASK_VALUE, + cap=attn_logits_soft_cap, ) if record_max_logits: # The native JAX splash attention implementation does not currently expose the softmax statistics @@ -1579,6 +1624,8 @@ def cudnn_flash_attention( and self.config.context_parallel_strategy == "ring" and self.config.context_parallel_load_balance ) + if using_context_parallelism and self.attention_type == AttentionType.CHUNK: + raise ValueError("Chunk attention is not supported with CUDNN context parallelism.") # Initialize default attention configuration sliding_window_size = None @@ -1806,6 +1853,7 @@ def apply_attention_dot( decoder_segment_ids: Array | None, model_mode: str = MODEL_MODE_TRAIN, previous_chunk: Any = None, + segment_positions: Array | None = None, bidirectional_mask: Any = None, sinks: Array | None = None, indexer_mask: Array | None = None, @@ -1871,6 +1919,7 @@ def apply_attention_dot( previous_chunk, bidirectional_mask, compressed_mask=compressed_mask, + segment_positions=segment_positions, ) if self.config.moba: @@ -2163,6 +2212,12 @@ def __call__( return prefill_unnormalized_output / prefill_exponentials_sum +def _load_balanced_q_sequence(shape: tuple[int, int], cp_size: int): + """Reorders query positions the same way as load-balanced input tokens.""" + arr = np.arange(shape[0]) + return max_utils.reorder_mask_load_balancing(arr, cp_size, 0) + + # pylint: disable=protected-access class LoadBalancedCausalMask(splash_attention_mask._ComputableMask): """Lazy causal mask, prevents the model from attending to future tokens. @@ -2193,12 +2248,6 @@ def causal_mask_function(q_ids, kv_ids): else: return q_ids + self.offset >= kv_ids - arr = np.arange(shape[0]) - # we reorder the mask to be load balanced following the same approach as - # used to reorder the input tokens - out = max_utils.reorder_mask_load_balancing(arr[None, :, None, None], cp_size, 1) - q_sequence = out[0, :, 0, 0] - mask_function = causal_mask_function super().__init__( @@ -2206,7 +2255,7 @@ def causal_mask_function(q_ids, kv_ids): mask_function=mask_function, shard_count=shard_count, ) - self.q_sequence = q_sequence + self.q_sequence = _load_balanced_q_sequence(shape, cp_size) def __eq__(self, other: object): if not isinstance(other, type(self)): @@ -2223,3 +2272,43 @@ def __hash__(self): self.q_sequence.tobytes() if self.q_sequence is not None else None, ) ) + + +class LoadBalancedLocalMask(splash_attention_mask.LocalMask): + """Lazy local mask with load-balanced query positions.""" + + def __init__( + self, + shape: tuple[int, int], + window_size: tuple[int | None, int | None], + offset: int, + shard_count: int = 1, + cp_size: int = -1, + ): + super().__init__( + shape=shape, + window_size=window_size, + offset=offset, + shard_count=shard_count, + ) + # LocalMask uses shard_count for mask-shard validation. cp_size is the + # context-parallel size used for the load-balanced query order. + self.q_sequence = _load_balanced_q_sequence(shape, cp_size) + + +class LoadBalancedChunkedCausalMask(ChunkedCausalMask): + """Lazy chunked mask with load-balanced query positions.""" + + def __init__( + self, + shape: tuple[int, int], + chunk_size: int, + cp_size: int, + shard_count: int = 1, + ): + super().__init__( + shape=shape, + chunk_size=chunk_size, + shard_count=shard_count, + ) + self.q_sequence = _load_balanced_q_sequence(shape, cp_size) diff --git a/src/maxtext/layers/decoders.py b/src/maxtext/layers/decoders.py index 0ea2913265..fc570e69bb 100644 --- a/src/maxtext/layers/decoders.py +++ b/src/maxtext/layers/decoders.py @@ -1056,6 +1056,8 @@ def __call__( bidirectional_mask_value, previous_chunk, slot, + kv_caches=kv_caches, + attention_metadata=attention_metadata, ) elif cfg.decoder_block == DecoderBlockType.GEMMA4: bidirectional_mask_value = multimodal_input.bidirectional_mask if multimodal_input is not None else None @@ -1068,6 +1070,8 @@ def __call__( bidirectional_mask_value, previous_chunk, slot, + kv_caches=kv_caches, + attention_metadata=attention_metadata, ) elif cfg.decoder_block == DecoderBlockType.DEEPSEEK4: y = self._apply_deepseek4_scanned_blocks( @@ -1303,6 +1307,8 @@ def _apply_gemma3_scanned_blocks( bidirectional_mask, previous_chunk, slot, + kv_caches=None, + attention_metadata=None, ): """Applies Gemma3 scanned decoder blocks, handling main scan and remainders.""" @@ -1316,27 +1322,43 @@ def _apply_gemma3_scanned_blocks( policy = self.get_remat_policy() RemattedGemma3Block = self.set_remat_policy([gemma3.Gemma3ScannableBlockToLinen], policy)[0] - layer_call_kwargs = {"bidirectional_mask": bidirectional_mask} layer_kwargs = {"num_of_layers": attention_pattern_length} # Apply the main scan over the full blocks if scan_length > 0: - broadcast_args = ( - decoder_segment_ids, - decoder_positions, - deterministic, - model_mode, + kv_cache_scanned = maxtext_utils.prepare_kv_caches_for_scan( + kv_caches, scan_length, attention_pattern_length, stack=True ) - y, _ = self.scan_decoder_layers( + + broadcast_args_spec = [ + (decoder_segment_ids, nn.broadcast), + (decoder_positions, nn.broadcast), + (deterministic, nn.broadcast), + (model_mode, nn.broadcast), + (slot, nn.broadcast), + (None, nn.broadcast), # page_state + (previous_chunk, nn.broadcast), + (bidirectional_mask, nn.broadcast), + (kv_cache_scanned, 0 if kv_caches is not None else nn.broadcast), + (attention_metadata, nn.broadcast), + ] + broadcast_args = tuple(arg for arg, _ in broadcast_args_spec) + in_axes_tuple = tuple(axis for _, axis in broadcast_args_spec) + + y, returned_kv_cache = self.scan_decoder_layers( cfg, RemattedGemma3Block, scan_length, "layers", mesh, - in_axes_tuple=(nn.broadcast,) * len(broadcast_args), + in_axes_tuple=in_axes_tuple, model_mode=self.model_mode, **layer_kwargs, - )(y, *broadcast_args, **layer_call_kwargs) + )(y, *broadcast_args) + + maxtext_utils.update_kv_caches_after_scan( + kv_caches, returned_kv_cache, scan_length, attention_pattern_length, stacked=True + ) # Apply any remaining layers that did not fit into a full scanned block num_remaining_layers = cfg.num_decoder_layers % attention_pattern_length @@ -1346,7 +1368,13 @@ def _apply_gemma3_scanned_blocks( layer = RemattedGemma3Block( config=cfg, mesh=mesh, quant=self.quant, model_mode=self.model_mode, name="layers_remainder", **rem_layer_kwargs ) # pytype: disable=wrong-keyword-args - y, _ = layer( + + remainder_kv = None + if kv_caches is not None: + start_idx = scan_length * attention_pattern_length + remainder_kv = tuple(kv_caches[start_idx : start_idx + num_remaining_layers]) + + y_and_kv = layer( y, decoder_segment_ids, decoder_positions, @@ -1354,8 +1382,23 @@ def _apply_gemma3_scanned_blocks( model_mode, previous_chunk=previous_chunk, slot=slot, - **layer_call_kwargs, + bidirectional_mask=bidirectional_mask, + kv_cache=remainder_kv, + attention_metadata=attention_metadata, ) + + if isinstance(y_and_kv, tuple): + y = y_and_kv[0] + updated_remainder_kv = y_and_kv[1] + else: + y = y_and_kv + updated_remainder_kv = None + + if kv_caches is not None and updated_remainder_kv is not None: + start_idx = scan_length * attention_pattern_length + for offset, updated_item in enumerate(updated_remainder_kv): + kv_caches[start_idx + offset] = updated_item + return y def _apply_gemma4_scanned_blocks( @@ -1368,6 +1411,8 @@ def _apply_gemma4_scanned_blocks( bidirectional_mask, previous_chunk, slot, + kv_caches=None, + attention_metadata=None, ): """Applies Gemma4 scanned decoder blocks, handling main scan and remainders.""" @@ -1379,28 +1424,32 @@ def _apply_gemma4_scanned_blocks( num_full_blocks = cfg.num_decoder_layers // block_pattern_len remainder_layers = cfg.num_decoder_layers % block_pattern_len - broadcast_args = ( - decoder_segment_ids, - decoder_positions, - deterministic, - model_mode, - ) - - # Pass slot/previous_chunk/bidirectional_mask by keyword only (via layer_call_kwargs), - # never positionally in broadcast_args: Gemma4DecoderLayer and Gemma4ScannableBlock - # declare slot and previous_chunk in swapped order, so positional passing misroutes them. - layer_call_kwargs = { - "slot": slot, - "previous_chunk": previous_chunk, - "bidirectional_mask": bidirectional_mask, - } - if num_full_blocks > 0: ScannableBlockToLinen = gemma4.Gemma4ScannableBlockToLinen policy = self.get_remat_policy() RemattedGemma4Block = self.set_remat_policy([ScannableBlockToLinen], policy)[0] + + kv_cache_scanned = maxtext_utils.prepare_kv_caches_for_scan( + kv_caches, num_full_blocks, block_pattern_len, stack=True + ) + + broadcast_args_spec = [ + (decoder_segment_ids, nn.broadcast), + (decoder_positions, nn.broadcast), + (deterministic, nn.broadcast), + (model_mode, nn.broadcast), + (slot, nn.broadcast), + (None, nn.broadcast), # page_state + (previous_chunk, nn.broadcast), + (bidirectional_mask, nn.broadcast), + (kv_cache_scanned, 0 if kv_caches is not None else nn.broadcast), + (attention_metadata, nn.broadcast), + ] + broadcast_args = tuple(arg for arg, _ in broadcast_args_spec) + in_axes_tuple = tuple(axis for _, axis in broadcast_args_spec) + # For a fully scanned block, apply it inside a nn.scan over the calculated number of full blocks - y, _ = nn.scan( + y, returned_kv_cache = nn.scan( RemattedGemma4Block, variable_axes={ "params": cfg.param_scan_axis, @@ -1410,7 +1459,7 @@ def _apply_gemma4_scanned_blocks( "_overwrite_with_gradient": 0, }, split_rngs={"params": True, "dropout": cfg.enable_dropout}, - in_axes=(nn.broadcast,) * len(broadcast_args), + in_axes=in_axes_tuple, length=num_full_blocks, metadata_params={ nn.PARTITION_NAME: "layers", @@ -1424,7 +1473,11 @@ def _apply_gemma4_scanned_blocks( num_of_layers=block_pattern_len, name="scanned_blocks", )( - y, *broadcast_args, **layer_call_kwargs + y, *broadcast_args + ) + + maxtext_utils.update_kv_caches_after_scan( + kv_caches, returned_kv_cache, num_full_blocks, block_pattern_len, stacked=True ) # Process any remaining layers that don't fit into a full scanned block @@ -1438,9 +1491,33 @@ def _apply_gemma4_scanned_blocks( attention_type=attention_type, layer_idx=layer_id, ) - y = layer(y, *broadcast_args, **layer_call_kwargs) - if cfg.scan_layers: - y = y[0] + kv_cache = kv_caches[layer_id] if kv_caches is not None else None + + # inputs, decoder_segment_ids, decoder_positions, deterministic, model_mode, + # previous_chunk, page_state, slot, bidirectional_mask, kv_cache, attention_metadata + remainder_args = ( + decoder_segment_ids, + decoder_positions, + deterministic, + model_mode, + previous_chunk, + None, # page_state + slot, + bidirectional_mask, + kv_cache, + attention_metadata, + ) + + y_and_kv = layer(y, *remainder_args) + if isinstance(y_and_kv, tuple): + y = y_and_kv[0] + new_kv = y_and_kv[1] + else: + y = y_and_kv + new_kv = None + + if kv_caches is not None and new_kv is not None: + kv_caches[layer_id] = new_kv return y diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py index fdea21981c..afa2799f9b 100644 --- a/src/maxtext/layers/moe.py +++ b/src/maxtext/layers/moe.py @@ -501,9 +501,7 @@ def __init__( kernel_out_axis = np.arange(1, 2) # Pad the MoE input weight kernels for GMM_v2 execution when configured. moe_intermediate_dim = ( - self.config.padded_base_moe_mlp_dim - if self.config.padded_base_moe_mlp_dim is not None - else self.intermediate_dim + self.config.padded_base_moe_mlp_dim if self.config.padded_base_moe_mlp_dim is not None else self.intermediate_dim ) if quantizations.in_serve_mode(self.quant): @@ -654,9 +652,6 @@ def get_tensor_parallelism_size(self): return size return self.mesh.shape.get(self._tensor_parallelism_name, 1) - def get_tensor_transpose_parallelism_size(self): - return self.mesh.shape.get("tensor_transpose", 1) - def get_context_autoregressive_parallelism_size(self): return self.mesh.shape.get("context_autoregressive", 1) @@ -1407,18 +1402,10 @@ def get_routed_moe_shardings(is_batch_sharded_by_expert, has_input_ids): else: batch_logical_axis = "decode_batch_moe" - if self.get_tensor_transpose_parallelism_size() > 1: - input_partition_pspec = self._logical_to_mesh_axes( - (batch_logical_axis, "activation_norm_length", "activation_embed") - ) - w0_bias_pspec = self._logical_to_mesh_axes(("exp", None)) - w1_bias_pspec = self._logical_to_mesh_axes(("exp", None)) - wo_bias_pspec = self._logical_to_mesh_axes(("exp", "activation_embed")) - else: - input_partition_pspec = self._logical_to_mesh_axes((batch_logical_axis, "activation_norm_length", None)) - w0_bias_pspec = self._logical_to_mesh_axes(("exp", "activation_mlp")) - w1_bias_pspec = self._logical_to_mesh_axes(("exp", "activation_mlp")) - wo_bias_pspec = self._logical_to_mesh_axes(("exp", "activation_embed")) + input_partition_pspec = self._logical_to_mesh_axes((batch_logical_axis, "activation_norm_length", None)) + w0_bias_pspec = self._logical_to_mesh_axes(("exp", "activation_mlp")) + w1_bias_pspec = self._logical_to_mesh_axes(("exp", "activation_mlp")) + wo_bias_pspec = self._logical_to_mesh_axes(("exp", "activation_embed")) gate_logits_pspec = self._logical_to_mesh_axes((batch_logical_axis, "activation_norm_length", None)) # NOTE: deepseek2 has a different pattern @@ -1444,18 +1431,18 @@ def get_routed_moe_shardings(is_batch_sharded_by_expert, has_input_ids): wo_pspec = self._logical_to_mesh_axes(self.wo_kernel_axes) else: # special sharding for dsv3 to remove overhead between gmm/AG - w0_pspec = self._logical_to_mesh_axes(("embed_tensor_transpose", None, "mlp_no_fsdp")) - w1_pspec = self._logical_to_mesh_axes(("embed_tensor_transpose", None, "mlp_no_fsdp")) - wo_pspec = self._logical_to_mesh_axes(("embed_tensor_transpose", "mlp_no_fsdp", None)) + w0_pspec = self._logical_to_mesh_axes((None, None, "mlp_no_fsdp")) + w1_pspec = self._logical_to_mesh_axes((None, None, "mlp_no_fsdp")) + wo_pspec = self._logical_to_mesh_axes((None, "mlp_no_fsdp", None)) elif self.config.use_2d_fsdp_sharding: - w0_pspec = self._logical_to_mesh_axes(("embed_tensor_transpose", "mlp_no_fsdp", None)) - w1_pspec = self._logical_to_mesh_axes(("embed_tensor_transpose", "mlp_no_fsdp", None)) - wo_pspec = self._logical_to_mesh_axes(("embed_tensor_transpose", "mlp_no_fsdp", None)) + w0_pspec = self._logical_to_mesh_axes((None, "mlp_no_fsdp", None)) + w1_pspec = self._logical_to_mesh_axes((None, "mlp_no_fsdp", None)) + wo_pspec = self._logical_to_mesh_axes((None, "mlp_no_fsdp", None)) else: # These are the main shardings used by default - they use funky rules to AG over FSDP. - w0_pspec = self._logical_to_mesh_axes(("exp", "embed_tensor_transpose", "mlp_no_fsdp")) - w1_pspec = self._logical_to_mesh_axes(("exp", "embed_tensor_transpose", "mlp_no_fsdp")) - wo_pspec = self._logical_to_mesh_axes(("exp", "mlp_no_fsdp", "embed_tensor_transpose")) + w0_pspec = self._logical_to_mesh_axes(("exp", None, "mlp_no_fsdp")) + w1_pspec = self._logical_to_mesh_axes(("exp", None, "mlp_no_fsdp")) + wo_pspec = self._logical_to_mesh_axes(("exp", "mlp_no_fsdp", None)) return ( batch_logical_axis, input_partition_pspec, @@ -1669,9 +1656,6 @@ def gmm_up(x, w0, w1, w0_bias, w1_bias, gmm_fn, weight_gather): out = gmm_fn(x, w_fused, tiling=wi_tile_size, weight_gather_axes=wi_gather_axes) n = out.shape[-1] // 2 layer_w0, layer_w1 = out[:, :n], out[:, n:] - if self.get_tensor_transpose_parallelism_size() > 1: - layer_w0 = jax.lax.psum(layer_w0, "tensor_transpose") - layer_w1 = jax.lax.psum(layer_w1, "tensor_transpose") if self.config.mlp_bias: layer_w0 = layer_w0 + w0_bias layer_w1 = layer_w1 + w1_bias @@ -1684,8 +1668,6 @@ def gmm_up(x, w0, w1, w0_bias, w1_bias, gmm_fn, weight_gather): tiling=wi_tile_size, weight_gather_axes=wi_gather_axes, ) - if self.get_tensor_transpose_parallelism_size() > 1: - layer_w0 = jax.lax.psum(layer_w0, "tensor_transpose") if self.config.mlp_bias: layer_w0 = layer_w0 + w0_bias layer_w0 = adc.checkpoint_name(adc.checkpoint_name(layer_w0, "mlpwi_0"), "moe_mlpwi_0") @@ -1696,8 +1678,6 @@ def gmm_up(x, w0, w1, w0_bias, w1_bias, gmm_fn, weight_gather): tiling=wi_tile_size, weight_gather_axes=wi_gather_axes, ) - if self.get_tensor_transpose_parallelism_size() > 1: - layer_w1 = jax.lax.psum(layer_w1, "tensor_transpose") if self.config.mlp_bias: layer_w1 = layer_w1 + w1_bias layer_w1 = adc.checkpoint_name(layer_w1, "moe_mlpwi_1") @@ -1884,11 +1864,11 @@ def sparse_matmul_route_and_compute( if self.config.moe_fsdp_use_two_stage_all_gather: # Unshard on fsdp axis - w0_kernel = self._maybe_shard_with_logical(w0_kernel, ("exp_with_fsdp", "embed_tensor_transpose", "mlp")) - w1_kernel = self._maybe_shard_with_logical(w1_kernel, ("exp_with_fsdp", "embed_tensor_transpose", "mlp")) + w0_kernel = self._maybe_shard_with_logical(w0_kernel, ("exp_with_fsdp", None, "mlp")) + w1_kernel = self._maybe_shard_with_logical(w1_kernel, ("exp_with_fsdp", None, "mlp")) # Unshard on fsdp_transpose axis - wo_kernel = self._maybe_shard_with_logical(wo_kernel, ("exp_with_fsdp", "mlp", "embed_tensor_transpose")) + wo_kernel = self._maybe_shard_with_logical(wo_kernel, ("exp_with_fsdp", "mlp", None)) # Make sure XLA does not optimize by combining above All-Gather to unshard # on FSDP axis and the subsequent unshard on fsdp_transpose axis @@ -1897,14 +1877,11 @@ def sparse_matmul_route_and_compute( wo_kernel = jax.lax.optimization_barrier(wo_kernel) # Unshard on both fsdp and fsdp_transpose transpose - w0_kernel = self._maybe_shard_with_logical(w0_kernel, ("exp_with_fsdp", "embed_tensor_transpose", "mlp_no_fsdp")) - w1_kernel = self._maybe_shard_with_logical(w1_kernel, ("exp_with_fsdp", "embed_tensor_transpose", "mlp_no_fsdp")) - wo_kernel = self._maybe_shard_with_logical(wo_kernel, ("exp_with_fsdp", "mlp_no_fsdp", "embed_tensor_transpose")) + w0_kernel = self._maybe_shard_with_logical(w0_kernel, ("exp_with_fsdp", None, "mlp_no_fsdp")) + w1_kernel = self._maybe_shard_with_logical(w1_kernel, ("exp_with_fsdp", None, "mlp_no_fsdp")) + wo_kernel = self._maybe_shard_with_logical(wo_kernel, ("exp_with_fsdp", "mlp_no_fsdp", None)) - if self.get_tensor_transpose_parallelism_size() > 1: - input_axes = (batch_logical_axis, "activation_norm_length", "activation_embed") - else: - input_axes = (batch_logical_axis, "activation_norm_length", None) + input_axes = (batch_logical_axis, "activation_norm_length", None) gate_logits_axes = (batch_logical_axis, "activation_norm_length", None) # NOTE: deepseek2 has a different pattern diff --git a/src/maxtext/layers/nnx_decoders.py b/src/maxtext/layers/nnx_decoders.py index 1a1ed65924..a7ce5822f9 100644 --- a/src/maxtext/layers/nnx_decoders.py +++ b/src/maxtext/layers/nnx_decoders.py @@ -653,7 +653,7 @@ def _init_scanned_gemma4(self, decoder_block_classes, rngs, mesh): RemattedGemma4Block = gemma4.Gemma4ScannableBlock if scan_length > 0: - self.layers = self._create_scanned_layers( + self.scanned_blocks = self._create_scanned_layers( RemattedGemma4Block, length=scan_length, metadata_axis_name="layers", @@ -1111,38 +1111,22 @@ def get_deepseek(): DecoderBlockType.GEMMA: [gemma.GemmaDecoderLayer], DecoderBlockType.GEMMA2: [gemma2.Gemma2DecoderLayer], DecoderBlockType.GEMMA3: [gemma3.Gemma3DecoderLayer], - DecoderBlockType.GEMMA4: get_scannable( - gemma4.Gemma4DecoderLayer, gemma4.Gemma4ScannableBlock - ), + DecoderBlockType.GEMMA4: get_scannable(gemma4.Gemma4DecoderLayer, gemma4.Gemma4ScannableBlock), DecoderBlockType.GEMMA4_SMALL: [gemma4_small.Gemma4SmallDecoderLayer], DecoderBlockType.GPT3: [gpt3.Gpt3DecoderLayer], DecoderBlockType.QWEN2: [qwen2.Qwen2DecoderLayer], DecoderBlockType.QWEN3: [qwen3.Qwen3DecoderLayer], DecoderBlockType.QWEN3_MOE: [qwen3.Qwen3MoeDecoderLayer], - DecoderBlockType.QWEN3_CUSTOM_MOE: [ - qwen3_custom.Qwen3CustomMoeDecoderLayer - ], + DecoderBlockType.QWEN3_CUSTOM_MOE: [qwen3_custom.Qwen3CustomMoeDecoderLayer], DecoderBlockType.SIMPLE: [simple_layer.SimpleDecoderLayer], DecoderBlockType.SIMPLE_MLP: [simple_layer.SimpleMlpDecoderLayer], DecoderBlockType.DEEPSEEK: get_deepseek(), - DecoderBlockType.DEEPSEEK4: get_scannable( - deepseek4.DeepSeek4DecoderLayer, deepseek4.DeepSeek4ScannableBlock - ), - DecoderBlockType.GPT_OSS: get_scannable( - gpt_oss.GptOssDecoderLayer, gpt_oss.GptOssScannableBlock - ), - DecoderBlockType.QWEN3_NEXT: get_scannable( - qwen3.Qwen3NextDecoderLayer, qwen3.Qwen3NextScannableBlock - ), - DecoderBlockType.QWEN3_5: get_scannable( - qwen3_5.Qwen3_5DecoderLayer, qwen3_5.Qwen3_5ScannableBlock - ), - DecoderBlockType.LLAMA4: get_scannable( - llama4.Llama4DecoderLayer, llama4.Llama4ScannableBlock - ), - DecoderBlockType.OLMO3: get_scannable( - olmo3.Olmo3DecoderLayer, olmo3.Olmo3ScannableBlock - ), + DecoderBlockType.DEEPSEEK4: get_scannable(deepseek4.DeepSeek4DecoderLayer, deepseek4.DeepSeek4ScannableBlock), + DecoderBlockType.GPT_OSS: get_scannable(gpt_oss.GptOssDecoderLayer, gpt_oss.GptOssScannableBlock), + DecoderBlockType.QWEN3_NEXT: get_scannable(qwen3.Qwen3NextDecoderLayer, qwen3.Qwen3NextScannableBlock), + DecoderBlockType.QWEN3_5: get_scannable(qwen3_5.Qwen3_5DecoderLayer, qwen3_5.Qwen3_5ScannableBlock), + DecoderBlockType.LLAMA4: get_scannable(llama4.Llama4DecoderLayer, llama4.Llama4ScannableBlock), + DecoderBlockType.OLMO3: get_scannable(olmo3.Olmo3DecoderLayer, olmo3.Olmo3ScannableBlock), } if cfg.decoder_block not in layer_map: @@ -1360,6 +1344,8 @@ def _apply_embedding( "llama4-17b-16e", "llama4-17b-128e", "qwen3-omni-30b-a3b", + "qwen3-vl-2b", + "qwen3-vl-4b", "qwen3.5-35b-a3b", "qwen3.5-397b-a17b", }: @@ -1373,7 +1359,7 @@ def _apply_embedding( raise ValueError(f"Unsupported model_name for multimodal: {cfg.model_name}") if video_embeddings is not None and cfg.use_multimodal: - if cfg.model_name in {"qwen3-omni-30b-a3b", "qwen3.5-35b-a3b", "qwen3.5-397b-a17b"}: + if cfg.model_name in {"qwen3-omni-30b-a3b", "qwen3-vl-2b", "qwen3-vl-4b", "qwen3.5-35b-a3b", "qwen3.5-397b-a17b"}: y = mm_utils.merge_mm_embeddings( text_embeddings=y, multimodal_embeddings=video_embeddings, @@ -1664,6 +1650,19 @@ def __call__( model_mode, logical_partition_spec=logical_partition_spec, ) + elif self.is_gemma4: + y = self._apply_gemma4_scanned_blocks( + y, + decoder_segment_ids, + decoder_positions, + deterministic, + model_mode, + bidirectional_mask, + previous_chunk, + slot, + kv_caches=kv_caches, + attention_metadata=attention_metadata, + ) else: # Standard pipeline run (non-DeepSeek, incl. Gemma4 — matches Linen decoders.py). # Gemma4 routes through the pipeline here; _apply_gemma4_scanned_blocks is @@ -1932,6 +1931,8 @@ def _apply_gemma3_scanned_blocks( bidirectional_mask, previous_chunk, slot, + kv_caches=None, + attention_metadata=None, ): """Applies Gemma3 scanned decoder blocks, handling main scan and remainders.""" @@ -1943,10 +1944,20 @@ def _apply_gemma3_scanned_blocks( layer_args = (decoder_segment_ids, decoder_positions, deterministic, model_mode) layer_kwargs = {"bidirectional_mask": bidirectional_mask} + if attention_metadata is not None: + layer_kwargs["attention_metadata"] = attention_metadata # Apply the main scan over the full blocks if scan_length > 0: - y, self.layers, _ = self._apply_layers_sequentially(self.layers, y, *layer_args, length=scan_length, **layer_kwargs) + grouped_kv_caches = maxtext_utils.prepare_kv_caches_for_scan( + kv_caches, scan_length, attention_pattern_length, stack=False + ) + y, self.layers, _ = self._apply_layers_sequentially( + self.layers, y, *layer_args, length=scan_length, kv_caches_stacked=grouped_kv_caches, **layer_kwargs + ) + maxtext_utils.update_kv_caches_after_scan( + kv_caches, grouped_kv_caches, scan_length, attention_pattern_length, stacked=False + ) # Apply any remaining layers that did not fit into a full scanned block num_remaining_layers = cfg.num_decoder_layers % attention_pattern_length @@ -1954,23 +1965,38 @@ def _apply_gemma3_scanned_blocks( policy = self.get_remat_policy() prevent_cse = maxtext_utils.should_prevent_cse_in_remat(cfg) - def pure_gemma_fn(graphdef, state_in, y_in): + remainder_kv = None + if kv_caches is not None: + start_idx = scan_length * attention_pattern_length + remainder_kv = tuple(kv_caches[start_idx : start_idx + num_remaining_layers]) + + def pure_gemma_fn(graphdef, state_in, y_in, kv_in): merged_layer = nnx.merge(graphdef, state_in) - out_y, _ = merged_layer( - y_in, - *layer_args, - previous_chunk=previous_chunk, - slot=slot, - **layer_kwargs, - ) - return out_y, nnx.state(merged_layer) + call_kwargs = dict(layer_kwargs) + if kv_in is not None: + call_kwargs["kv_cache"] = kv_in + call_kwargs["previous_chunk"] = previous_chunk + call_kwargs["slot"] = slot + out_res = merged_layer(y_in, *layer_args, **call_kwargs) + if isinstance(out_res, tuple): + out_y = out_res[0] + out_kv = out_res[1] if len(out_res) > 1 else None + else: + out_y = out_res + out_kv = None + return out_y, out_kv, nnx.state(merged_layer) checkpointed_gemma_fn = jax.checkpoint(pure_gemma_fn, policy=policy, prevent_cse=prevent_cse) graphdef, state = nnx.split(self.layers_remainder) - y, new_state = checkpointed_gemma_fn(graphdef, state, y) + y, updated_remainder_kv, new_state = checkpointed_gemma_fn(graphdef, state, y, remainder_kv) nnx.update(self.layers_remainder, new_state) + if kv_caches is not None and updated_remainder_kv is not None: + start_idx = scan_length * attention_pattern_length + for offset, updated_item in enumerate(updated_remainder_kv): + kv_caches[start_idx + offset] = updated_item + return y def _apply_gemma4_scanned_blocks( @@ -1983,6 +2009,8 @@ def _apply_gemma4_scanned_blocks( bidirectional_mask, previous_chunk, slot, + kv_caches=None, + attention_metadata=None, ): """Applies Gemma4 scanned decoder blocks, handling main scan and remainders.""" @@ -1994,10 +2022,20 @@ def _apply_gemma4_scanned_blocks( layer_args = (decoder_segment_ids, decoder_positions, deterministic, model_mode) layer_kwargs = {"bidirectional_mask": bidirectional_mask, "slot": slot, "previous_chunk": previous_chunk} + if attention_metadata is not None: + layer_kwargs["attention_metadata"] = attention_metadata # Apply the main scan over the full blocks if scan_length > 0: - y, self.layers, _ = self._apply_layers_sequentially(self.layers, y, *layer_args, length=scan_length, **layer_kwargs) + grouped_kv_caches = maxtext_utils.prepare_kv_caches_for_scan( + kv_caches, scan_length, attention_pattern_length, stack=False + ) + y, self.scanned_blocks, _ = self._apply_layers_sequentially( + self.scanned_blocks, y, *layer_args, length=scan_length, kv_caches_stacked=grouped_kv_caches, **layer_kwargs + ) + maxtext_utils.update_kv_caches_after_scan( + kv_caches, grouped_kv_caches, scan_length, attention_pattern_length, stacked=False + ) # Apply any remaining layers that did not fit into a full scanned block num_remaining_layers = cfg.num_decoder_layers % attention_pattern_length @@ -2005,21 +2043,38 @@ def _apply_gemma4_scanned_blocks( policy = self.get_remat_policy() prevent_cse = maxtext_utils.should_prevent_cse_in_remat(cfg) - def pure_gemma_fn(graphdef, state_in, y_in): + remainder_kv = None + if kv_caches is not None: + start_idx = scan_length * attention_pattern_length + remainder_kv = tuple(kv_caches[start_idx : start_idx + num_remaining_layers]) + + def pure_gemma_fn(graphdef, state_in, y_in, kv_in): merged_layer = nnx.merge(graphdef, state_in) - out_y, _ = merged_layer( - y_in, - *layer_args, - **layer_kwargs, - ) - return out_y, nnx.state(merged_layer) + call_kwargs = dict(layer_kwargs) + if kv_in is not None: + call_kwargs["kv_cache"] = kv_in + call_kwargs["previous_chunk"] = previous_chunk + call_kwargs["slot"] = slot + out_res = merged_layer(y_in, *layer_args, **call_kwargs) + if isinstance(out_res, tuple): + out_y = out_res[0] + out_kv = out_res[1] if len(out_res) > 1 else None + else: + out_y = out_res + out_kv = None + return out_y, out_kv, nnx.state(merged_layer) checkpointed_gemma_fn = jax.checkpoint(pure_gemma_fn, policy=policy, prevent_cse=prevent_cse) graphdef, state = nnx.split(self.layers_remainder) - y, new_state = checkpointed_gemma_fn(graphdef, state, y) + y, updated_remainder_kv, new_state = checkpointed_gemma_fn(graphdef, state, y, remainder_kv) nnx.update(self.layers_remainder, new_state) + if kv_caches is not None and updated_remainder_kv is not None: + start_idx = scan_length * attention_pattern_length + for offset, updated_item in enumerate(updated_remainder_kv): + kv_caches[start_idx + offset] = updated_item + return y def _apply_gemma4_small_layers( diff --git a/src/maxtext/layers/nnx_wrappers.py b/src/maxtext/layers/nnx_wrappers.py index a1d4ef64fe..2f5756023a 100644 --- a/src/maxtext/layers/nnx_wrappers.py +++ b/src/maxtext/layers/nnx_wrappers.py @@ -237,7 +237,7 @@ def __getattr__(self, name: str): maybe_method = getattr(type(self.to_nnx__module), name, None) if callable(maybe_method): method = partial(self.__call__, method=maybe_method) - method.__self__ = self + method.__self__ = self # pyrefly: ignore[missing-attribute] return method return super().__getattribute__(name) @@ -361,14 +361,14 @@ def wrapped(*args, **kwargs): if not linen.module._context.module_stack: # pylint: disable=W0212 return call_fn(*args, **kwargs) nn_module = linen.module._context.module_stack[-1] # pylint: disable=W0212 - old_path = nn_module.path + old_path = nn_module.path # pyrefly: ignore[missing-attribute] # We modify the path of the current nn module in place. This is a little # bit hacky but should be good as a temporary solution. - nn_module.scope.path += (name,) + nn_module.scope.path += (name,) # pyrefly: ignore[missing-attribute] try: return call_fn(*args, **kwargs) finally: - nn_module.scope.path = old_path + nn_module.scope.path = old_path # pyrefly: ignore[missing-attribute] return wrapped @@ -556,7 +556,7 @@ def __getattr__(self, name: str): maybe_method = getattr(self.nnx_class, name, None) if callable(maybe_method): method = partial(self.__call__, nnx_method=maybe_method) - method.__self__ = self + method.__self__ = self # pyrefly: ignore[missing-attribute] return method return super().__getattribute__(name) @@ -701,12 +701,12 @@ class ToLinenPartial(ToLinen): def __init_subclass__(cls, **kwargs): super().__init_subclass__(**kwargs) - cls.__init__ = __init__ + cls.__init__ = __init__ # pyrefly: ignore[bad-assignment] ToLinenPartial.__name__ = class_name ToLinenPartial.__qualname__ = class_name - ToLinenPartial.__init__ = __init__ + ToLinenPartial.__init__ = __init__ # pyrefly: ignore[bad-assignment] ToLinenPartial.module_class = base_nnx_class return ToLinenPartial diff --git a/src/maxtext/layers/quantizations.py b/src/maxtext/layers/quantizations.py index 7174cf763a..d016b66b07 100644 --- a/src/maxtext/layers/quantizations.py +++ b/src/maxtext/layers/quantizations.py @@ -103,22 +103,22 @@ def _rhs_axis_metadata_wrapper( if len(x.shape) == 1: return nn.with_logical_partitioning((lambda: x), tuple(None for _ in mesh_axes))() - mesh_axes = list(mesh_axes) + mesh_axes = list(mesh_axes) # pyrefly: ignore[bad-assignment] if is_tiled: # tile_map is a mapping between original rank and a list of new, tiled rank. if len(mesh_axes) < len(tile_map): - mesh_axes = [None] * (len(tile_map) - len(mesh_axes)) + mesh_axes + mesh_axes = [None] * (len(tile_map) - len(mesh_axes)) + mesh_axes # pyrefly: ignore[unsupported-operation] new_mesh_axes = [None] * len(x.shape) for orig_rank, new_rank in tile_map.items(): assert new_rank assert len(new_rank) <= 2 new_mesh_axes[new_rank[-1]] = mesh_axes[orig_rank] - mesh_axes = new_mesh_axes + mesh_axes = new_mesh_axes # pyrefly: ignore[bad-assignment] if mesh_axes is not None and len(mesh_axes) > 0: for no_shard_idx in no_sharding_axis: if no_shard_idx < len(mesh_axes): - mesh_axes[no_shard_idx] = None + mesh_axes[no_shard_idx] = None # pyrefly: ignore[unsupported-operation] return nn.with_logical_partitioning((lambda: x), mesh_axes)() @@ -137,13 +137,13 @@ def _get_mixed_precision_cfg(self): is_tiled = False tiling_fn = None # pylint: disable=protected-access - module_path = "/".join(nn.module._context.module_stack[-1].path) + module_path = "/".join(nn.module._context.module_stack[-1].path) # pyrefly: ignore[missing-attribute] tile_size = -1 - for layer_name_re, layer_quant_dg in self.quant_dg.items(): + for layer_name_re, layer_quant_dg in self.quant_dg.items(): # pyrefly: ignore[missing-attribute] if re.fullmatch(layer_name_re, module_path): quant_dg, tile_size = layer_quant_dg if quant_dg is None: - quant_dg, tile_size = self.quant_dg[DEFAULT] + quant_dg, tile_size = self.quant_dg[DEFAULT] # pyrefly: ignore[bad-index] if tile_size != -1: is_tiled = True tiling_fn = functools.partial(_tiling_fn, tile_size=tile_size) @@ -272,7 +272,7 @@ def __call__( out_sharding=None, ) -> jax.Array: def custom_dot_general(*args, **kwargs): - return dot_general_qt.dot_general_qt(*args[:3], self.config) + return dot_general_qt.dot_general_qt(*args[:3], self.config) # pyrefly: ignore[bad-argument-count] with jax.disable_jit(): return jnp.einsum( @@ -392,7 +392,7 @@ def _get_int8_quant_config(config): if config.quantization_local_shard_count != 0: drhs_bits = 8 drhs_accumulator_dtype = jnp.int32 - drhs_local_aqt = aqt_config.LocalAqt(contraction_axis_shard_count=config.quantization_local_shard_count) + drhs_local_aqt = aqt_config.LocalAqt(contraction_axis_shard_count=config.quantization_local_shard_count) # pyrefly: ignore[unexpected-keyword] return aqt_config.config_v3( fwd_bits=8, dlhs_bits=8, @@ -430,30 +430,30 @@ def _build_const_scale_config( The AQT dot general config with constant scale config. """ if cst_bound_config.fwd_lhs_bound is not None: - aqt_dg.fwd.dg_quantizer.lhs.calibration = functools.partial( + aqt_dg.fwd.dg_quantizer.lhs.calibration = functools.partial( # pyrefly: ignore[missing-attribute] calibration.ConstantCalibration, bound=cst_bound_config.fwd_lhs_bound ) if cst_bound_config.fwd_rhs_bound is not None: - aqt_dg.fwd.dg_quantizer.rhs.calibration = functools.partial( + aqt_dg.fwd.dg_quantizer.rhs.calibration = functools.partial( # pyrefly: ignore[missing-attribute] calibration.ConstantCalibration, bound=cst_bound_config.fwd_rhs_bound ) if cst_bound_config.dlhs_lhs_bound: - aqt_dg.dlhs.dg_quantizer.lhs.calibration = functools.partial( + aqt_dg.dlhs.dg_quantizer.lhs.calibration = functools.partial( # pyrefly: ignore[missing-attribute] calibration.ConstantCalibration, bound=cst_bound_config.dlhs_lhs_bound ) if cst_bound_config.dlhs_rhs_bound is not None: - aqt_dg.dlhs.dg_quantizer.rhs.calibration = functools.partial( + aqt_dg.dlhs.dg_quantizer.rhs.calibration = functools.partial( # pyrefly: ignore[missing-attribute] calibration.ConstantCalibration, bound=cst_bound_config.dlhs_rhs_bound ) if cst_bound_config.drhs_lhs_bound is not None: - aqt_dg.drhs.dg_quantizer.lhs.calibration = functools.partial( + aqt_dg.drhs.dg_quantizer.lhs.calibration = functools.partial( # pyrefly: ignore[missing-attribute] calibration.ConstantCalibration, bound=cst_bound_config.drhs_lhs_bound ) if cst_bound_config.drhs_rhs_bound is not None: - aqt_dg.drhs.dg_quantizer.rhs.calibration = functools.partial( + aqt_dg.drhs.dg_quantizer.rhs.calibration = functools.partial( # pyrefly: ignore[missing-attribute] calibration.ConstantCalibration, bound=cst_bound_config.drhs_rhs_bound ) @@ -484,17 +484,17 @@ def _build_per_tensor_config( The AQT dot general config with per tensor config. """ if per_tensor_scales.fwd_lhs: - aqt_dg.fwd.dg_quantizer.lhs.calib_shared_axes = "per_tensor" + aqt_dg.fwd.dg_quantizer.lhs.calib_shared_axes = "per_tensor" # pyrefly: ignore[missing-attribute] if per_tensor_scales.fwd_rhs: - aqt_dg.fwd.dg_quantizer.rhs.calib_shared_axes = "per_tensor" + aqt_dg.fwd.dg_quantizer.rhs.calib_shared_axes = "per_tensor" # pyrefly: ignore[missing-attribute] if per_tensor_scales.dlhs_lhs: - aqt_dg.dlhs.dg_quantizer.lhs.calib_shared_axes = "per_tensor" + aqt_dg.dlhs.dg_quantizer.lhs.calib_shared_axes = "per_tensor" # pyrefly: ignore[missing-attribute] if per_tensor_scales.dlhs_rhs: - aqt_dg.dlhs.dg_quantizer.rhs.calib_shared_axes = "per_tensor" + aqt_dg.dlhs.dg_quantizer.rhs.calib_shared_axes = "per_tensor" # pyrefly: ignore[missing-attribute] if per_tensor_scales.drhs_lhs: - aqt_dg.drhs.dg_quantizer.lhs.calib_shared_axes = "per_tensor" + aqt_dg.drhs.dg_quantizer.lhs.calib_shared_axes = "per_tensor" # pyrefly: ignore[missing-attribute] if per_tensor_scales.drhs_rhs: - aqt_dg.drhs.dg_quantizer.rhs.calib_shared_axes = "per_tensor" + aqt_dg.drhs.dg_quantizer.rhs.calib_shared_axes = "per_tensor" # pyrefly: ignore[missing-attribute] return aqt_dg @@ -565,9 +565,9 @@ def _dot_general_make(quant_cfg): rhs_scale = quant_cfg[_W_SCALE] aqt_dg = aqt_config.dot_general_make(lhs_bits=lhs_bits, rhs_bits=rhs_bits) if lhs_scale < 1.0: - aqt_dg.fwd.dg_quantizer.lhs.calibration = functools.partial(calibration.AbsMaxCalibration, scale=lhs_scale) + aqt_dg.fwd.dg_quantizer.lhs.calibration = functools.partial(calibration.AbsMaxCalibration, scale=lhs_scale) # pyrefly: ignore[missing-attribute] if rhs_scale < 1.0: - aqt_dg.fwd.dg_quantizer.rhs.calibration = functools.partial(calibration.AbsMaxCalibration, scale=rhs_scale) + aqt_dg.fwd.dg_quantizer.rhs.calibration = functools.partial(calibration.AbsMaxCalibration, scale=rhs_scale) # pyrefly: ignore[missing-attribute] return aqt_dg @@ -728,7 +728,7 @@ def _apply_linen_module_in_nnx(linen_module_cls, op_id, *args, **kwargs): if is_nnx: attr_name = f"_qwix_fp8_gpu_{op_id}" - if not hasattr(parent, attr_name): + if not hasattr(parent, attr_name): # pyrefly: ignore[unbound-name] rngs = getattr(parent, "qwix_rngs", None) if rngs is None: parent_rngs = getattr(parent, "rngs", None) @@ -951,6 +951,8 @@ def __init__(self, config): self._recipe = TransformerEngineQuantization._get_recipe(config.quantization) + self._perform_collective_gemm = config.use_te_comm_gemm_overlap + def __hash__(self): return hash((self.quant_mode, self._recipe)) @@ -1001,11 +1003,13 @@ def _wrap(self, f, name=None): 2. Wraps the given function in a Flax linen module. This module does not store any Flax parameters but can store Flax variables for quantizers if required by the recipe. - 3. When the wrapper is called, it provides an additional argument to the given function `f`, - 'generate_quantizer_set' as the first argument. 'generate_quantizer_set' is a function that - can be called to generate a TransformerEngine/JAX quantizer set object used in - TransformerEngine/JAX APIs. 'generate_quantizer_set' will generate quantizers based on the - recipe of this TransformerEngineQuantizer object. + 3. When the wrapper is called, it provides two additional arguments to the given function `f`, + 'generate_quantizer_set' as the first argument and 'generate_collective_op_set' as the second argument. + 'generate_quantizer_set' is a function that can be called to generate a TransformerEngine/JAX quantizer + set object used in TransformerEngine/JAX APIs based on the recipe of this TransformerEngineQuantizer + object. Similarly, 'generate_collective_op_set' is a function that can be called to generate a + TransformerEngine/JAX collective operation set object used in TransformerEngine/JAX APIs based + the kernel's mesh axes. Args: f: The function to wrap. The first argument must be 'generate_quantizer_set'. @@ -1016,6 +1020,7 @@ def _wrap(self, f, name=None): """ import transformer_engine.jax # pylint: disable=import-outside-toplevel # pytype: disable=import-error + import transformer_engine.jax.cpp_extensions as tex # pylint: disable=import-outside-toplevel # pytype: disable=import-error from transformer_engine.common import recipe # pylint: disable=import-outside-toplevel # pytype: disable=import-error default_recipe = self._recipe @@ -1038,12 +1043,23 @@ def generate_quantizer_set( variable_collection=OVERWRITE_WITH_GRADIENT, quantization_checkpoint_name="quantization", fp8_recipe=fp8_recipe, - n_groups=n_groups, + n_groups=n_groups, # pyrefly: ignore[bad-argument-type] ) + def generate_collective_op_set(self, mesh_axes: Tuple[str, ...] = ()): + """Inspect the kernel's mesh axes to determine the type of collective operation to use for collective GEMM.""" + + if len(mesh_axes) >= 1: + if mesh_axes[0] == "embed" and mesh_axes[-1] == "mlp": + return tex.CollectiveOpSet.create(tex.CollectiveOp.ALL_GATHER) + elif mesh_axes[0] == "mlp" and mesh_axes[-1] == "embed": + return tex.CollectiveOpSet.create(tex.CollectiveOp.REDUCE_SCATTER) + + return tex.noop_collective_op_set + @nn.compact def __call__(self, *args, **kwargs): - return f(self.generate_quantizer_set, *args, **kwargs) + return f(self.generate_quantizer_set, self.generate_collective_op_set, *args, **kwargs) TEWrapper.__name__ = f"TEWrapper_{name if name else f.__name__}" @@ -1052,17 +1068,22 @@ def __call__(self, *args, **kwargs): def dot_general_cls(self, mesh_axes: Tuple[str, ...] = ()): """Placeholder for dot_general implementation in subclasses.""" import transformer_engine.jax # pylint: disable=import-outside-toplevel # pytype: disable=import-error + import transformer_engine.jax.cpp_extensions as tex # pylint: disable=import-outside-toplevel # pytype: disable=import-error - def te_dot_general(generate_quantizer_set, x, kernel, dims, **kwargs): + def te_dot_general(generate_quantizer_set, generate_collective_op_set, x, kernel, dims, **kwargs): contracting_dims, batch_dims = dims assert batch_dims == ((), ()), "Batch dimensions must be empty for TransformerEngine dot." quantizer_set = generate_quantizer_set() + collective_op_set = ( + generate_collective_op_set(mesh_axes) if self._perform_collective_gemm else tex.noop_collective_op_set + ) return transformer_engine.jax.dense.dense( x, kernel, contracting_dims=contracting_dims, quantizer_set=quantizer_set, + collective_op_set=collective_op_set, ) return self._wrap(te_dot_general, "dot_general") diff --git a/src/maxtext/models/deepseek.py b/src/maxtext/models/deepseek.py index d3a72b31bf..3ec5c32f1f 100644 --- a/src/maxtext/models/deepseek.py +++ b/src/maxtext/models/deepseek.py @@ -303,9 +303,9 @@ def self_attention_with_norm_op( return hidden_states, intermediate_inputs def engram_op(self, x, decoder_input_tokens): - normed_x = self.engram_layer_norm(x) + normed_x = self.engram_layer_norm(x) # pyrefly: ignore[not-callable] hash_ids = self.ngram_hash_mapping(decoder_input_tokens)[self.layer_idx] - return self.engram(normed_x, hash_ids) + return self.engram(normed_x, hash_ids) # pyrefly: ignore[not-callable] class DeepSeekDenseLayer(DeepSeekGenericLayer): diff --git a/src/maxtext/models/deepseek_batchsplit_fp8.py b/src/maxtext/models/deepseek_batchsplit_fp8.py index cef7c0646f..0b6410c9e5 100644 --- a/src/maxtext/models/deepseek_batchsplit_fp8.py +++ b/src/maxtext/models/deepseek_batchsplit_fp8.py @@ -1004,7 +1004,7 @@ def gmm( gating_pspec, linear_pspec = moe_lib.get_batchsplit_init_kernel_axes() w0_pspec = nn.logical_to_mesh_axes(gating_pspec) wo_pspec = nn.logical_to_mesh_axes(linear_pspec) - ignored_axes = ("expert", "tensor", "tensor_transpose") + ignored_axes = ("expert", "tensor") def get_active_sharding_axes(pspec_dim_axes, tensor_dim_index): if pspec_dim_axes is None: diff --git a/src/maxtext/models/gemma3.py b/src/maxtext/models/gemma3.py index 2d2b36a431..c33272d277 100644 --- a/src/maxtext/models/gemma3.py +++ b/src/maxtext/models/gemma3.py @@ -259,8 +259,6 @@ def update_cache(cache, val): stacked_kv_cache = jax.tree_util.tree_map(update_cache, stacked_kv_cache, kv_cache) return (layer_output, stacked_kv_cache, layer_idx + 1), None - elif cfg.scan_layers: - return layer_output, None else: return layer_output, kv_cache @@ -324,6 +322,8 @@ def __call__( page_state=None, previous_chunk=None, bidirectional_mask=None, + kv_cache=None, + attention_metadata=None, ): cfg = self.config @@ -331,8 +331,10 @@ def __call__( inputs = checkpoint_name(inputs, "decoder_layer_input") y = inputs + updated_kvs = [] for layer_id in range(self.num_of_layers): - y = getattr(self, f"layers_{layer_id}")( + current_kv = kv_cache[layer_id] if kv_cache is not None else None + y, new_kv = getattr(self, f"layers_{layer_id}")( y, decoder_segment_ids, decoder_positions, @@ -341,10 +343,15 @@ def __call__( previous_chunk=previous_chunk, slot=slot, bidirectional_mask=bidirectional_mask, + kv_cache=current_kv, + attention_metadata=attention_metadata, ) - if cfg.scan_layers: - y = y[0] - if cfg.scan_layers: + if kv_cache is not None: + updated_kvs.append(new_kv) + + if kv_cache is not None: + return y, tuple(updated_kvs) + elif cfg.scan_layers: return y, None else: return y diff --git a/src/maxtext/models/gemma4.py b/src/maxtext/models/gemma4.py index 8cf927629c..dbf8efc4de 100644 --- a/src/maxtext/models/gemma4.py +++ b/src/maxtext/models/gemma4.py @@ -398,8 +398,6 @@ def update_cache(cache, val): stacked_kv_cache = jax.tree_util.tree_map(update_cache, stacked_kv_cache, kv_cache) return (layer_output, stacked_kv_cache, layer_idx + 1), None - elif cfg.scan_layers: - return layer_output, None else: return layer_output, kv_cache @@ -464,13 +462,18 @@ def __call__( page_state=None, previous_chunk=None, bidirectional_mask=None, + kv_cache=None, + attention_metadata=None, ): + cfg = self.config inputs = nn.with_logical_constraint(inputs, ("activation_batch", "activation_norm_length", "activation_embed")) inputs = checkpoint_name(inputs, "decoder_layer_input") y = inputs + updated_kvs = [] for layer_id in range(self.num_of_layers): - y, _ = getattr(self, f"layers_{layer_id}")( + current_kv = kv_cache[layer_id] if kv_cache is not None else None + y, new_kv = getattr(self, f"layers_{layer_id}")( y, decoder_segment_ids, decoder_positions, @@ -479,9 +482,18 @@ def __call__( previous_chunk=previous_chunk, slot=slot, bidirectional_mask=bidirectional_mask, + kv_cache=current_kv, + attention_metadata=attention_metadata, ) + if kv_cache is not None: + updated_kvs.append(new_kv) - return y, None + if kv_cache is not None: + return y, tuple(updated_kvs) + elif cfg.scan_layers: + return y, None + else: + return y Gemma4ScannableBlockToLinen = nnx_wrappers.to_linen_class( diff --git a/src/maxtext/models/gpt3.py b/src/maxtext/models/gpt3.py index 7cf8fbd773..b09523e453 100644 --- a/src/maxtext/models/gpt3.py +++ b/src/maxtext/models/gpt3.py @@ -330,7 +330,7 @@ def init_kv_caches(self, inputs_kv_shape: tuple[int, ...]): ) def update_kv_caches(self, key, value, decoder_segment_ids, model_mode, previous_chunk): - prefill_kv_cache, ar_kv_cache = self.KVCache_0( + prefill_kv_cache, ar_kv_cache = self.KVCache_0( # pyrefly: ignore[not-callable] key=key, value=value, decoder_segment_ids=decoder_segment_ids, diff --git a/src/maxtext/models/llama2.py b/src/maxtext/models/llama2.py index eeb54da934..fa0f8b85e5 100644 --- a/src/maxtext/models/llama2.py +++ b/src/maxtext/models/llama2.py @@ -218,8 +218,8 @@ def update_cache(cache, val): return cache.at[layer_idx].set(val) return cache - stacked_kv_cache = jax.tree_util.tree_map(update_cache, stacked_kv_cache, kv_cache) - return (layer_output, stacked_kv_cache, layer_idx + 1), None + stacked_kv_cache = jax.tree_util.tree_map(update_cache, stacked_kv_cache, kv_cache) # pyrefly: ignore[unbound-name] + return (layer_output, stacked_kv_cache, layer_idx + 1), None # pyrefly: ignore[unbound-name] elif cfg.scan_layers: return layer_output, None else: diff --git a/src/maxtext/models/llama4.py b/src/maxtext/models/llama4.py index 26fd4d322d..65a22c41e2 100644 --- a/src/maxtext/models/llama4.py +++ b/src/maxtext/models/llama4.py @@ -50,7 +50,7 @@ class Llama4UnfoldConvolution(nnx.Module): config: Config containing model parameters """ - def __init__(self, config: Config, *, rngs: nnx.Rngs = None): + def __init__(self, config: Config, *, rngs: nnx.Rngs = None): # pyrefly: ignore[bad-function-definition] self.config = config self.rngs = rngs self.vit_unfold_linear = linears.DenseGeneral( @@ -123,7 +123,7 @@ class Llama4VisionMLP(nnx.Module): config: Config containing model parameters """ - def __init__(self, config: Config, *, rngs: nnx.Rngs = None): + def __init__(self, config: Config, *, rngs: nnx.Rngs = None): # pyrefly: ignore[bad-function-definition] self.config = config self.rngs = rngs self.vit_encoder_layer_mlp_fc1 = linears.DenseGeneral( @@ -157,7 +157,7 @@ class Llama4VisionMLP2(nnx.Module): config: Config containing model parameters """ - def __init__(self, config: Config, *, rngs: nnx.Rngs = None): + def __init__(self, config: Config, *, rngs: nnx.Rngs = None): # pyrefly: ignore[bad-function-definition] self.config = config self.rngs = rngs self.vit_pixel_shuffle_mlp_fc1 = linears.DenseGeneral( @@ -196,7 +196,7 @@ class Llama4VisionPixelShuffleMLP(nnx.Module): config: Config containing model parameters """ - def __init__(self, config: Config, *, rngs: nnx.Rngs = None): + def __init__(self, config: Config, *, rngs: nnx.Rngs = None): # pyrefly: ignore[bad-function-definition] self.config = config self.rngs = rngs self.pixel_shuffle_ratio = self.config.pixel_shuffle_ratio_for_vit @@ -221,7 +221,7 @@ class Llama4MultiModalProjector(nnx.Module): config: Config containing model parameters """ - def __init__(self, config: Config, mesh: Mesh, *, rngs: nnx.Rngs = None): + def __init__(self, config: Config, mesh: Mesh, *, rngs: nnx.Rngs = None): # pyrefly: ignore[bad-function-definition] self.config = config self.mesh = mesh self.rngs = rngs @@ -515,8 +515,8 @@ def update_cache(cache, val): return cache.at[layer_idx].set(val) return cache - stacked_kv_cache = jax.tree_util.tree_map(update_cache, stacked_kv_cache, kv_cache) - return (layer_output, stacked_kv_cache, layer_idx + 1), None + stacked_kv_cache = jax.tree_util.tree_map(update_cache, stacked_kv_cache, kv_cache) # pyrefly: ignore[unbound-name] + return (layer_output, stacked_kv_cache, layer_idx + 1), None # pyrefly: ignore[unbound-name] elif cfg.scan_layers: return layer_output, None else: @@ -623,7 +623,7 @@ def __call__( class Llama4VisionEncoderLayer(nnx.Module): """Transformer encoder layer for Llama4 vision model.""" - def __init__(self, config: Config, mesh: Mesh, *, rngs: nnx.Rngs = None): + def __init__(self, config: Config, mesh: Mesh, *, rngs: nnx.Rngs = None): # pyrefly: ignore[bad-function-definition] self.config = config self.mesh = mesh self.rngs = rngs @@ -701,7 +701,7 @@ class Llama4VisionEncoder(nnx.Module): mesh: Mesh, JAX device mesh (used for sharding) """ - def __init__(self, config: Config, mesh: Mesh, *, rngs: nnx.Rngs = None): + def __init__(self, config: Config, mesh: Mesh, *, rngs: nnx.Rngs = None): # pyrefly: ignore[bad-function-definition] self.config = config self.mesh = mesh self.rngs = rngs @@ -733,7 +733,7 @@ class Llama4VisionModel(nnx.Module): mesh: Mesh, JAX device mesh (used for sharding) """ - def __init__(self, config: Config, mesh: Mesh, *, rngs: nnx.Rngs = None): + def __init__(self, config: Config, mesh: Mesh, *, rngs: nnx.Rngs = None): # pyrefly: ignore[bad-function-definition] self.config = config self.mesh = mesh self.rngs = rngs diff --git a/src/maxtext/models/mistral.py b/src/maxtext/models/mistral.py index fa7f0956d4..78fbcfa1d7 100644 --- a/src/maxtext/models/mistral.py +++ b/src/maxtext/models/mistral.py @@ -192,8 +192,8 @@ def update_cache(cache, val): return cache.at[layer_idx].set(val) return cache - stacked_kv_cache = jax.tree_util.tree_map(update_cache, stacked_kv_cache, kv_cache) - return (layer_output, stacked_kv_cache, layer_idx + 1), None + stacked_kv_cache = jax.tree_util.tree_map(update_cache, stacked_kv_cache, kv_cache) # pyrefly: ignore[unbound-name] + return (layer_output, stacked_kv_cache, layer_idx + 1), None # pyrefly: ignore[unbound-name] elif cfg.scan_layers: return layer_output, None else: diff --git a/src/maxtext/models/models.py b/src/maxtext/models/models.py index ac908c0f96..881a1687b5 100644 --- a/src/maxtext/models/models.py +++ b/src/maxtext/models/models.py @@ -102,9 +102,9 @@ def setup(self): self.mtp_block = multi_token_prediction_block_as_linen( config=self.config, mesh=self.mesh, - transformer_layer_module=mtp_layer_nnx, - decoder=self.decoder, - rngs=self.make_rng("mtp_block"), + transformer_layer_module=mtp_layer_nnx, # pyrefly: ignore[bad-argument-type] + decoder=self.decoder, # pyrefly: ignore[bad-argument-type] + rngs=self.make_rng("mtp_block"), # pyrefly: ignore[bad-argument-type] ) def logits_from_hidden_states_for_vocab_tiling(self, hidden_states, deterministic, model_mode): @@ -163,7 +163,7 @@ def __call__( deepstack_visual_embeds = None if self.config.use_multimodal and encoder_images is not None: - image_embeddings, deepstack_visual_embeds = self.vision_encoder( + image_embeddings, deepstack_visual_embeds = self.vision_encoder( # pyrefly: ignore[not-callable] input_images=encoder_images, deterministic=not enable_dropout ) bidirectional_mask_image = mm_processor.get_bidirectional_mask_vision( @@ -171,7 +171,7 @@ def __call__( ) if self.config.use_multimodal and encoder_videos is not None: - video_embeddings, deepstack_visual_embeds = self.vision_encoder( + video_embeddings, deepstack_visual_embeds = self.vision_encoder( # pyrefly: ignore[not-callable] input_images=encoder_videos, deterministic=not enable_dropout ) bidirectional_mask_video = mm_processor.get_bidirectional_mask_vision( @@ -384,7 +384,7 @@ def __init__( dummy_attention_metadata = None if not cfg.pure_nnx_decoder: - self.decoder.lazy_init( + self.decoder.lazy_init( # pyrefly: ignore[missing-attribute] shared_embedding=self.token_embedder, decoder_input_tokens=dummy_decoder_input_tokens, decoder_positions=dummy_decoder_positions, @@ -490,7 +490,7 @@ def __call__( audio_embeddings = None deepstack_visual_embeds = None if self.config.use_multimodal and encoder_images is not None: - image_embeddings, deepstack_visual_embeds = self.vision_encoder( + image_embeddings, deepstack_visual_embeds = self.vision_encoder( # pyrefly: ignore[not-callable] input_images=encoder_images, deterministic=not enable_dropout ) bidirectional_mask_image = mm_processor.get_bidirectional_mask_vision( @@ -498,7 +498,7 @@ def __call__( ) if self.config.use_multimodal and encoder_videos is not None: - video_embeddings, deepstack_visual_embeds = self.vision_encoder( + video_embeddings, deepstack_visual_embeds = self.vision_encoder( # pyrefly: ignore[not-callable] input_images=encoder_videos, deterministic=not enable_dropout ) bidirectional_mask_video = mm_processor.get_bidirectional_mask_vision( @@ -563,7 +563,7 @@ def __call__( kv_caches=kv_caches, attention_metadata=attention_metadata, deepstack_visual_embeds=deepstack_visual_embeds, - mutable=mutable_collections, + mutable=mutable_collections, # pyrefly: ignore[unexpected-keyword] ) # pytype: disable=wrong-keyword-args # If we are initializing the model AND MTP is enabled, we must create diff --git a/src/maxtext/multimodal/processor.py b/src/maxtext/multimodal/processor.py index 381695915a..f15617bf10 100644 --- a/src/maxtext/multimodal/processor.py +++ b/src/maxtext/multimodal/processor.py @@ -149,15 +149,15 @@ def prepare_text_for_image_fusion(tokens, config, processor_output=None): if config.model_name in ["gemma3-4b", "gemma3-12b", "gemma3-27b"]: from maxtext.multimodal.processor_gemma3 import add_extra_tokens_for_images_gemma3 # pylint: disable=import-outside-toplevel - return add_extra_tokens_for_images_gemma3(tokens, max_num_images=processor_output.num_images) + return add_extra_tokens_for_images_gemma3(tokens, max_num_images=processor_output.num_images) # pyrefly: ignore[missing-attribute] elif config.model_name in ["gemma4-26b", "gemma4-31b", "gemma4-e2b", "gemma4-e4b"]: from maxtext.multimodal.processor_gemma4 import add_extra_tokens_for_images_gemma4 # pylint: disable=import-outside-toplevel - return add_extra_tokens_for_images_gemma4(tokens, max_num_images=processor_output.num_images) + return add_extra_tokens_for_images_gemma4(tokens, max_num_images=processor_output.num_images) # pyrefly: ignore[missing-attribute] elif config.model_name in ["llama4-17b-16e", "llama4-17b-128e"]: from maxtext.multimodal.processor_llama4 import add_extra_tokens_for_images_llama4 # pylint: disable=import-outside-toplevel - return add_extra_tokens_for_images_llama4(tokens, processor_output) + return add_extra_tokens_for_images_llama4(tokens, processor_output) # pyrefly: ignore[bad-argument-type] elif config.model_name in ["qwen3-omni-30b-a3b", "qwen3-vl-2b", "qwen3-vl-4b", "qwen3.5-35b-a3b", "qwen3.5-397b-a17b"]: from maxtext.multimodal.processor_qwen3_omni import add_extra_tokens_for_qwen3_omni # pylint: disable=import-outside-toplevel diff --git a/src/maxtext/multimodal/processor_gemma3.py b/src/maxtext/multimodal/processor_gemma3.py index 4b444fe935..7e059edd46 100644 --- a/src/maxtext/multimodal/processor_gemma3.py +++ b/src/maxtext/multimodal/processor_gemma3.py @@ -86,7 +86,7 @@ def preprocess_mm_data_gemma3(images): def get_image_offsets_gemma3(processor_output: mm_utils.PreprocessorOutput | None): """Get the increase in total token count after inserting image token placeholders""" has_images = processor_output is not None and processor_output.pixel_values is not None - num_images = processor_output.pixel_values.shape[0] if has_images else 1 + num_images = processor_output.pixel_values.shape[0] if has_images else 1 # pyrefly: ignore[missing-attribute] return ( GEMMA_NUM_TOKENS_PER_MEDIA - 1 ) * num_images # -1 because is already present in the input tokens. diff --git a/src/maxtext/multimodal/processor_gemma4.py b/src/maxtext/multimodal/processor_gemma4.py index 5671c750ac..e261eb8223 100644 --- a/src/maxtext/multimodal/processor_gemma4.py +++ b/src/maxtext/multimodal/processor_gemma4.py @@ -78,7 +78,7 @@ def preprocess_mm_data_gemma4(images): def get_image_offsets_gemma4(processor_output: mm_utils.PreprocessorOutput | None): """Gets the increase in total token count after inserting image token placeholders.""" has_images = processor_output is not None and processor_output.pixel_values is not None - num_images = processor_output.pixel_values.shape[0] if has_images else 1 + num_images = processor_output.pixel_values.shape[0] if has_images else 1 # pyrefly: ignore[missing-attribute] # Calculate soft tokens taking 3x3 pooling into account num_patches = (GEMMA4_IMAGE_HEIGHT // GEMMA4_PATCH_SIZE) * (GEMMA4_IMAGE_WIDTH // GEMMA4_PATCH_SIZE) diff --git a/src/maxtext/multimodal/processor_qwen3_omni.py b/src/maxtext/multimodal/processor_qwen3_omni.py index 787dd82011..08a59f1506 100644 --- a/src/maxtext/multimodal/processor_qwen3_omni.py +++ b/src/maxtext/multimodal/processor_qwen3_omni.py @@ -167,7 +167,7 @@ def maybe_pad_video_values_to_max_grid( if video_values.ndim != 5: raise ValueError(f"video_values must have shape (batch, channels, time, height, width), got {video_values.shape}.") - max_t, max_h, max_w = (int(dim) for dim in max_grid) + max_t, max_h, max_w = (int(dim) for dim in max_grid) # pyrefly: ignore[bad-argument-type] actual_t, actual_h, actual_w = (int(dim) for dim in video_grid_thw[0]) if actual_t > max_t or actual_h > max_h or actual_w > max_w: raise ValueError( @@ -199,7 +199,13 @@ def maybe_pad_video_values_to_max_grid( def smart_resize( - height: int, width: int, factor: int = 28, min_pixels: int = 56 * 56, max_pixels: int = 14 * 14 * 4 * 1280 + height: int, + width: int, + factor: int = 28, + min_pixels: int = 56 * 56, + max_pixels: int = 14 * 14 * 4 * 1280, + max_h_pixels: int | None = None, + max_w_pixels: int | None = None, ): """Rescales the image so that the following conditions are met: @@ -209,6 +215,10 @@ def smart_resize( 3. The aspect ratio of the image is maintained as closely as possible. + 4. (Optional) Neither dimension exceeds max_h_pixels / max_w_pixels. When either + cap is exceeded, both dimensions are scaled down proportionally and re-aligned + to 'factor'. + """ if max(height, width) / min(height, width) > MAX_RATIO: raise ValueError( @@ -224,6 +234,13 @@ def smart_resize( beta = math.sqrt(min_pixels / (height * width)) h_bar = math.ceil(height * beta / factor) * factor w_bar = math.ceil(width * beta / factor) * factor + # Apply per-dimension pixel caps, scaling down proportionally if either is exceeded. + if (max_h_pixels is not None and h_bar > max_h_pixels) or (max_w_pixels is not None and w_bar > max_w_pixels): + cap_h = max_h_pixels if max_h_pixels is not None else h_bar + cap_w = max_w_pixels if max_w_pixels is not None else w_bar + scale = min(cap_h / h_bar, cap_w / w_bar) + h_bar = min(max(factor, round(h_bar * scale / factor) * factor), (cap_h // factor) * factor) + w_bar = min(max(factor, round(w_bar * scale / factor) * factor), (cap_w // factor) * factor) return h_bar, w_bar @@ -348,6 +365,7 @@ def calculate_video_frame_range( # Validate frame order if start_frame >= end_frame: raise ValueError( + # pyrefly: ignore[unbound-name] f"Invalid time range: Start frame {start_frame} (at {video_start_clamped if video_start is not None else 0}s) " f"exceeds end frame {end_frame} (at {video_end_clamped if video_end is not None else max_duration}s). " f"Video duration: {max_duration:.2f}s ({total_frames} frames @ {video_fps}fps)" @@ -461,6 +479,10 @@ def preprocess_video(video, config): nframes, channel, height, width = video.shape max_pixels = max(min(VIDEO_MAX_PIXELS, VIDEO_TOTAL_PIXELS / nframes * FRAME_FACTOR), int(VIDEO_MIN_PIXELS * 1.05)) + max_grid_h = getattr(config, "video_max_grid_h", None) + max_grid_w = getattr(config, "video_max_grid_w", None) + max_h_pixels = int(max_grid_h) * patch_size if max_grid_h is not None else None + max_w_pixels = int(max_grid_w) * patch_size if max_grid_w is not None else None resized_height_1, resized_width_1 = smart_resize( height, width, @@ -488,6 +510,8 @@ def preprocess_video(video, config): factor=patch_size * merge_size, min_pixels=VIDEO_MIN_PIXELS, max_pixels=VIDEO_MAX_PIXELS, + max_h_pixels=max_h_pixels, + max_w_pixels=max_w_pixels, ) # Second resize - process each channel separately to preserve float values @@ -780,7 +804,7 @@ def add_extra_tokens_for_qwen3_omni(tokens, config, processor_output): new_tokens.append(qwen_tokens.audio_pad) audio_data_idx += 1 - new_tokens.append(qwen_tokens.audio_pad) + new_tokens.append(qwen_tokens.audio_end) new_tokens.append(qwen_tokens.vision_end) video_idx += 1 @@ -1112,7 +1136,7 @@ def get_rope_index( # Process modality-specific content # Audio Only if min_ed == ed_audio_start: - audio_len = _get_feat_extract_output_lengths(audio_lengths[audio_idx]).item() + audio_len = _get_feat_extract_output_lengths(audio_lengths[audio_idx]).item() # pyrefly: ignore[unsupported-operation] audio_pos = np.arange(audio_len).reshape(1, -1).repeat(3, axis=0) + st_idx llm_pos_ids_list.append(audio_pos) @@ -1122,45 +1146,45 @@ def get_rope_index( # Image Only elif min_ed == ed_vision_start and input_tokens[ed_vision_start + 1] == qwen_tokens.image_pad: - grid_t = image_grid_thw[image_idx, 0].item() - grid_hs = image_grid_thw[:, 1] - grid_ws = image_grid_thw[:, 2] + grid_t = image_grid_thw[image_idx, 0].item() # pyrefly: ignore[unsupported-operation] + grid_hs = image_grid_thw[:, 1] # pyrefly: ignore[unsupported-operation] + grid_ws = image_grid_thw[:, 2] # pyrefly: ignore[unsupported-operation] t_index = np.arange(grid_t, dtype=np.float32) * 1 * position_id_per_seconds - image_pos = get_llm_pos_ids_for_vision(st_idx, image_idx, spatial_merge_size, t_index, grid_hs, grid_ws) + image_pos = get_llm_pos_ids_for_vision(st_idx, image_idx, spatial_merge_size, t_index, grid_hs, grid_ws) # pyrefly: ignore[bad-argument-type] llm_pos_ids_list.append(image_pos) - image_len = int(np.prod(image_grid_thw[image_idx]).item() // (spatial_merge_size**2)) + image_len = int(np.prod(image_grid_thw[image_idx]).item() // (spatial_merge_size**2)) # pyrefly: ignore[unsupported-operation] st += int(text_len + bos_len + image_len + eos_len) image_idx += 1 remain_images -= 1 # Video Only elif min_ed == ed_vision_start and input_tokens[ed_vision_start + 1] == qwen_tokens.video_pad: - grid_t = video_grid_thw[video_idx, 0].item() - grid_hs = video_grid_thw[:, 1] - grid_ws = video_grid_thw[:, 2] - t_index = np.arange(grid_t, dtype=np.float32) * second_per_grids[video_idx].item() * position_id_per_seconds + grid_t = video_grid_thw[video_idx, 0].item() # pyrefly: ignore[unsupported-operation] + grid_hs = video_grid_thw[:, 1] # pyrefly: ignore[unsupported-operation] + grid_ws = video_grid_thw[:, 2] # pyrefly: ignore[unsupported-operation] + t_index = np.arange(grid_t, dtype=np.float32) * second_per_grids[video_idx].item() * position_id_per_seconds # pyrefly: ignore[unsupported-operation] - video_pos = get_llm_pos_ids_for_vision(st_idx, video_idx, spatial_merge_size, t_index, grid_hs, grid_ws) + video_pos = get_llm_pos_ids_for_vision(st_idx, video_idx, spatial_merge_size, t_index, grid_hs, grid_ws) # pyrefly: ignore[bad-argument-type] llm_pos_ids_list.append(video_pos) - video_len = int(np.prod(video_grid_thw[video_idx]).item() // (spatial_merge_size**2)) + video_len = int(np.prod(video_grid_thw[video_idx]).item() // (spatial_merge_size**2)) # pyrefly: ignore[unsupported-operation] st += int(text_len + bos_len + video_len + eos_len) video_idx += 1 remain_videos -= 1 # Audio in Video (interleaved) elif min_ed == ed_vision_start and ed_vision_start + 1 == ed_audio_start: - audio_len = _get_feat_extract_output_lengths(audio_lengths[audio_idx]).item() + audio_len = _get_feat_extract_output_lengths(audio_lengths[audio_idx]).item() # pyrefly: ignore[unsupported-operation] audio_llm_pos_ids = np.arange(audio_len).reshape(1, -1).repeat(3, axis=0) + st_idx - grid_t = video_grid_thw[video_idx, 0].item() - grid_hs = video_grid_thw[:, 1] - grid_ws = video_grid_thw[:, 2] - t_index = np.arange(grid_t, dtype=np.float32) * second_per_grids[video_idx].item() * position_id_per_seconds + grid_t = video_grid_thw[video_idx, 0].item() # pyrefly: ignore[unsupported-operation] + grid_hs = video_grid_thw[:, 1] # pyrefly: ignore[unsupported-operation] + grid_ws = video_grid_thw[:, 2] # pyrefly: ignore[unsupported-operation] + t_index = np.arange(grid_t, dtype=np.float32) * second_per_grids[video_idx].item() * position_id_per_seconds # pyrefly: ignore[unsupported-operation] - video_llm_pos_ids = get_llm_pos_ids_for_vision(st_idx, video_idx, spatial_merge_size, t_index, grid_hs, grid_ws) + video_llm_pos_ids = get_llm_pos_ids_for_vision(st_idx, video_idx, spatial_merge_size, t_index, grid_hs, grid_ws) # pyrefly: ignore[bad-argument-type] # Interleave audio and video based on temporal ordering video_data_index = 0 @@ -1179,7 +1203,7 @@ def get_rope_index( if audio_data_index < audio_llm_pos_ids.shape[1]: llm_pos_ids_list.append(audio_llm_pos_ids[:, audio_data_index:]) - video_len = int(np.prod(video_grid_thw[video_idx]).item() // (spatial_merge_size**2)) + video_len = int(np.prod(video_grid_thw[video_idx]).item() // (spatial_merge_size**2)) # pyrefly: ignore[unsupported-operation] st += int(text_len + bos_len + audio_len + video_len + eos_len) audio_idx += 1 diff --git a/src/maxtext/multimodal/utils.py b/src/maxtext/multimodal/utils.py index 65b5670fc1..a068cc6fa9 100644 --- a/src/maxtext/multimodal/utils.py +++ b/src/maxtext/multimodal/utils.py @@ -188,7 +188,7 @@ def _merge_mm_embeddings_inner( # Find positions in the text sequence to place the multimodal embeddings. # The `size` argument ensures a fixed shape for JIT compilation. - target_pos = jnp.nonzero(mask, size=multimodal_embeddings.shape[0]) + target_pos = jnp.nonzero(mask, size=multimodal_embeddings.shape[0]) # pyrefly: ignore[bad-argument-type] target_pos = target_pos[0] # jnp.nonzero returns a tuple of arrays # Save the embedding at the first position. @@ -452,7 +452,7 @@ def spectrogram( # center pad the waveform if center: padding = [(int(frame_length // 2), int(frame_length // 2))] - waveform = np.pad(waveform, padding, mode=pad_mode) + waveform = np.pad(waveform, padding, mode=pad_mode) # pyrefly: ignore[no-matching-overload] # promote to float64, since np.fft uses float64 internally waveform = waveform.astype(np.float64) @@ -545,7 +545,7 @@ def hertz_to_mel(freq: Union[float, np.ndarray], mel_scale: str = "htk") -> Unio if isinstance(freq, np.ndarray): log_region = freq >= min_log_hertz - mels[log_region] = min_log_mel + np.log(freq[log_region] / min_log_hertz) * logstep + mels[log_region] = min_log_mel + np.log(freq[log_region] / min_log_hertz) * logstep # pyrefly: ignore[unsupported-operation] elif freq >= min_log_hertz: mels = min_log_mel + np.log(freq / min_log_hertz) * logstep @@ -603,7 +603,7 @@ def mel_to_hertz(mels: Union[float, np.ndarray], mel_scale: str = "htk") -> Unio if isinstance(mels, np.ndarray): log_region = mels >= min_log_mel - freq[log_region] = min_log_hertz * np.exp(logstep * (mels[log_region] - min_log_mel)) + freq[log_region] = min_log_hertz * np.exp(logstep * (mels[log_region] - min_log_mel)) # pyrefly: ignore[unsupported-operation] elif mels >= min_log_mel: freq = min_log_hertz * np.exp(logstep * (mels - min_log_mel)) @@ -688,11 +688,11 @@ def mel_filter_bank( # frequencies of FFT bins in Hz fft_freqs = np.linspace(0, sampling_rate // 2, num_frequency_bins) - mel_filters = _create_triangular_filter_bank(fft_freqs, filter_freqs) + mel_filters = _create_triangular_filter_bank(fft_freqs, filter_freqs) # pyrefly: ignore[bad-argument-type] if norm is not None and norm == "slaney": # Slaney-style mel is scaled to be approx constant energy per channel - enorm = 2.0 / (filter_freqs[2 : num_mel_filters + 2] - filter_freqs[:num_mel_filters]) + enorm = 2.0 / (filter_freqs[2 : num_mel_filters + 2] - filter_freqs[:num_mel_filters]) # pyrefly: ignore[bad-index] mel_filters *= np.expand_dims(enorm, 0) if (mel_filters.max(axis=0) == 0.0).any(): diff --git a/src/maxtext/trainers/diloco/diloco.py b/src/maxtext/trainers/diloco/diloco.py index 1b7253c41d..48accf4779 100644 --- a/src/maxtext/trainers/diloco/diloco.py +++ b/src/maxtext/trainers/diloco/diloco.py @@ -157,11 +157,11 @@ def add_diloco_dim(x): # for Linen under abstract_state.params. if config.pure_nnx: _, model_params, _ = nnx.split(abstract_state.model, nnx.Param, ...) - model_params = model_params.to_pure_dict() + model_params = model_params.to_pure_dict() # pyrefly: ignore[missing-attribute] _, model_params_sharding, _ = nnx.split( state_mesh_shardings.model, nnx.Param, ... ) - model_params_sharding = model_params_sharding.to_pure_dict() + model_params_sharding = model_params_sharding.to_pure_dict() # pyrefly: ignore[missing-attribute] else: model_params = abstract_state.params model_params_sharding = state_mesh_shardings.params @@ -175,7 +175,7 @@ def add_diloco_dim(x): inner_state=inner_state, params=model_params, outer_opt_state=outer_opt_state, - step=abstract_step, + step=abstract_step, # pyrefly: ignore[bad-argument-type] ) # Build shardings @@ -190,7 +190,7 @@ def add_diloco_dim(x): inner_state=inner_state_shardings, params=model_params_sharding, outer_opt_state=outer_opt_state_sharding, - step=None, + step=None, # pyrefly: ignore[bad-argument-type] ) return diloco_state, diloco_state_shardings, inner_state_shardings @@ -221,7 +221,7 @@ def init_diloco_state() -> tuple[DiLoCoTrainState, PyTree]: # for Linen under state.params. if config.pure_nnx: _, outer_params, _ = nnx.split(state.model, nnx.Param, ...) - outer_params = outer_params.to_pure_dict() + outer_params = outer_params.to_pure_dict() # pyrefly: ignore[missing-attribute] else: outer_params = state.params outer_opt_state = outer_optimizer.init(outer_params) @@ -269,7 +269,7 @@ def synchronize(state): _, inner_model_params, _ = nnx.split( state.inner_state.model, nnx.Param, ... ) - inner_model_params = inner_model_params.to_pure_dict() + inner_model_params = inner_model_params.to_pure_dict() # pyrefly: ignore[missing-attribute] else: inner_model_params = state.inner_state.params model_delta = jax.tree.map(lambda x, y: y - x, inner_model_params, broadcast_outer_params) diff --git a/src/maxtext/trainers/post_train/distillation/distillation_utils.py b/src/maxtext/trainers/post_train/distillation/distillation_utils.py index f063cdb23a..1946c2f2e8 100644 --- a/src/maxtext/trainers/post_train/distillation/distillation_utils.py +++ b/src/maxtext/trainers/post_train/distillation/distillation_utils.py @@ -508,7 +508,7 @@ def compute_loss( log_t_p_T_sparse = jax.nn.log_softmax(t_logits / temperature, axis=-1) # 2. Gather Student unnormalized logits at the Teacher's exact Top-K indices - s_logits_sparse = jnp.take_along_axis(s_logits, teacher_output.top_k_indices, axis=-1) + s_logits_sparse = jnp.take_along_axis(s_logits, teacher_output.top_k_indices, axis=-1) # pyrefly: ignore[bad-argument-type] # 3. Normalize Student probabilities only over the exact same Top-K subset log_s_T_sparse = jax.nn.log_softmax(s_logits_sparse / temperature, axis=-1) @@ -558,7 +558,7 @@ def compute_loss( s_features_sliced = s_features_sliced.astype(jnp.float32) t_features_sliced = t_features_sliced.astype(jnp.float32) - feature_loss = beta_feature * self.feature_loss_fn(s_features_sliced, t_features_sliced, mask) + feature_loss = beta_feature * self.feature_loss_fn(s_features_sliced, t_features_sliced, mask) # pyrefly: ignore[not-callable] total_loss = base_logit_loss + feature_loss @@ -744,7 +744,7 @@ def save( grain_iters_to_save.append((local_iter, process_index, process_count_total)) # Use GrainCheckpointSave wrapper - cp_save_args["iter"] = GrainCheckpointSave(item=grain_iters_to_save) + cp_save_args["iter"] = GrainCheckpointSave(item=grain_iters_to_save) # pyrefly: ignore[bad-assignment] return self._checkpoint_manager.save( step, @@ -753,7 +753,7 @@ def save( force=force, ) - def maybe_restore( + def maybe_restore( # pyrefly: ignore[bad-override] self, model: Any, optimizer: Any = None, @@ -772,7 +772,7 @@ def maybe_restore( target_model = getattr(model, "student_model", model) step, _ = super().maybe_restore( - model=target_model, + model=target_model, # pyrefly: ignore[bad-argument-type] optimizer=optimizer, restore_only_lora_params=restore_only_lora_params, ) diff --git a/src/maxtext/trainers/post_train/distillation/train_distill.py b/src/maxtext/trainers/post_train/distillation/train_distill.py index 15197d1af4..952811bd37 100644 --- a/src/maxtext/trainers/post_train/distillation/train_distill.py +++ b/src/maxtext/trainers/post_train/distillation/train_distill.py @@ -203,7 +203,7 @@ def call_student(self, *args, **kwargs): return self.student_model(*args, **kwargs) def call_teacher(self, *args, **kwargs): - return jax.lax.stop_gradient(self.teacher_model(*args, **kwargs)) + return jax.lax.stop_gradient(self.teacher_model(*args, **kwargs)) # pyrefly: ignore[not-callable] class MaxTextDistillationTrainer(peft_trainer.PeftTrainer): @@ -232,7 +232,7 @@ def __init__( super().__init__(model=model, optimizer=dummy_optimizer, training_config=training_config, **kwargs) self.strategy = strategy - self.checkpoint_manager: distillation_utils.MaxTextCheckpointManager = None + self.checkpoint_manager: distillation_utils.MaxTextCheckpointManager = None # pyrefly: ignore[bad-assignment] # Per-step per-device TFLOPs (constants for the run): student fwd+bwd + teacher fwd-only. ( @@ -385,11 +385,12 @@ def _log_metrics(self, loss, step=None, additional_metrics=None, **kwargs): } ) for name, value in tflops_metrics.items(): - self.metrics_logger.log(self.metrics_prefix, name, value, self._mode, step) + self.metrics_logger.log(self.metrics_prefix, name, value, self._mode, step) # pyrefly: ignore[missing-attribute] # Console summary — keep it tight; everything else is in TensorBoard. if tflops_per_sec is not None and self._mode == metrics_logger.Mode.TRAIN: max_logging.log( + # pyrefly: ignore[unsupported-operation] f"step {step} | step_time={step_time_delta:.2f}s | " f"TFLOPs/s/device: {tflops_per_sec:.2f} " f"(student={self._tflops_student / step_time_delta:.2f}, " @@ -438,9 +439,9 @@ def _post_process_train_step(self, aux: dict[str, tuple[jax.Array, jax.Array]]) for name, value in aux.items(): if name not in self._buffered_train_metrics.additional_metrics: - self._buffered_train_metrics.additional_metrics[name] = ([], distillation_utils.weighted_mean) + self._buffered_train_metrics.additional_metrics[name] = ([], distillation_utils.weighted_mean) # pyrefly: ignore[unsupported-operation] - self._buffered_train_metrics.additional_metrics[name][0].append(value) + self._buffered_train_metrics.additional_metrics[name][0].append(value) # pyrefly: ignore[bad-argument-type] # Compact per-step summary: only the metrics that change run-to-run, and # only those that are nonzero (so MoE / feature-distill terms hide when off). diff --git a/src/maxtext/trainers/post_train/rl/math_verify_pool.py b/src/maxtext/trainers/post_train/rl/math_verify_pool.py index 2eeadf1747..c0968f184e 100644 --- a/src/maxtext/trainers/post_train/rl/math_verify_pool.py +++ b/src/maxtext/trainers/post_train/rl/math_verify_pool.py @@ -226,9 +226,9 @@ def math_verify_pool( cpu_count = multiprocessing.cpu_count() if num_procs is None: - num_procs = min(_DEFAULT_MAX_PROCS, len(items), cpu_count) + num_procs = min(_DEFAULT_MAX_PROCS, cpu_count) else: - num_procs = max(1, min(num_procs, len(items), cpu_count)) + num_procs = max(1, min(num_procs, cpu_count)) cnt = 0 pool = _get_pool(num_procs) diff --git a/src/maxtext/trainers/post_train/rl/train_rl.py b/src/maxtext/trainers/post_train/rl/train_rl.py index 9270fb5033..19f714f7ae 100644 --- a/src/maxtext/trainers/post_train/rl/train_rl.py +++ b/src/maxtext/trainers/post_train/rl/train_rl.py @@ -46,7 +46,7 @@ from __future__ import annotations import contextlib from functools import wraps -from typing import Any, Optional, Sequence +from typing import Any, Callable, Optional, Sequence import datasets import grain @@ -339,7 +339,7 @@ def _use_raw_prompt(x): dataset_size = int(trainer_config.num_batches * trainer_config.batch_size * trainer_config.train_fraction) train_dataset = train_dataset[:dataset_size] train_dataset = train_dataset.repeat(trainer_config.num_epoch) - train_dataset = train_dataset.to_iter_dataset().batch(trainer_config.batch_size) + train_dataset = train_dataset.to_iter_dataset().batch(trainer_config.batch_size, drop_remainder=True) if trainer_config.num_test_batches > 0: # eval_batch_size = -1 (default) → use trainer_config.batch_size (legacy @@ -358,11 +358,36 @@ def _use_raw_prompt(x): test_dataset = test_dataset[ trainer_config.test_batch_start_index : trainer_config.num_test_batches * eval_batch_size_for_eval ] - test_dataset = test_dataset.to_iter_dataset().batch(eval_batch_size_for_eval) + test_dataset = test_dataset.to_iter_dataset().batch(eval_batch_size_for_eval, drop_remainder=True) return train_dataset, test_dataset +def build_reward_fns(trainer_config: Any, make_reward_fn: Callable) -> list: + """Build the reward-function stack for the RL trainer. + + `reward_functions_path` is a filesystem path to a Python file and + `reward_functions` is a comma-separated list of function names to import from + it. When both are set, the built-in stack is REPLACED entirely by the + user-provided callables (so users have full control over their reward stack). + Otherwise the default + (`match_format_exactly`, `match_format_approximately`, `check_numbers`) stack + is used. Every reward function is wrapped via `make_reward_fn`. + """ + custom_rewards_path = getattr(trainer_config, "reward_functions_path", "") or "" + custom_rewards_names = getattr(trainer_config, "reward_functions", "") or "" + if custom_rewards_path and custom_rewards_names: + names = [n.strip() for n in custom_rewards_names.split(",") if n.strip()] + reward_fns = [make_reward_fn(utils_rl.load_custom_callable(custom_rewards_path, n)) for n in names] + max_logging.log(f"reward_fns: using {len(reward_fns)} custom reward function(s) {names} from {custom_rewards_path}") + return reward_fns + return [ + make_reward_fn(utils_rl.match_format_exactly), + make_reward_fn(utils_rl.match_format_approximately), + make_reward_fn(utils_rl.check_numbers), + ] + + def create_rl_components( trainer_config, sampler_config, @@ -526,11 +551,11 @@ def _reward_fn(**kwargs): return _reward_fn - reward_fns = [ # type: ignore - make_reward_fn(utils_rl.match_format_exactly), - make_reward_fn(utils_rl.match_format_approximately), - make_reward_fn(utils_rl.check_numbers), - ] + # Optional user-provided reward functions: when `reward_functions_path` and + # `reward_functions` are both set the built-in stack is replaced entirely by + # the user-provided callables. Each function must accept `prompts`, + # `completions`, `tmvp_config`, and `**kwargs` and return a list of floats. + reward_fns = build_reward_fns(trainer_config, make_reward_fn) # Create RL trainer max_logging.log("Setting up RL trainer...") @@ -555,7 +580,6 @@ def _reward_fn(**kwargs): max_concurrency=trainer_config.rl.max_concurrency, off_policy_steps=trainer_config.rl.off_policy_steps, system_prompt=trainer_config.rl.system_prompt, - degenerate_group_masking=trainer_config.rl.degenerate_group_masking, epsilon_high=trainer_config.rl.epsilon_high, ) # Instantiate the custom MaxText chat parser diff --git a/src/maxtext/trainers/post_train/sft/train_sft.py b/src/maxtext/trainers/post_train/sft/train_sft.py index 4e185e93a1..9002caa16c 100644 --- a/src/maxtext/trainers/post_train/sft/train_sft.py +++ b/src/maxtext/trainers/post_train/sft/train_sft.py @@ -262,8 +262,8 @@ def setup_trainer_state(mt_config, goodput_recorder=None): # Provide rules context so 'norm' is translated to mesh axes during maybe_restore with nn_partitioning.axis_rules(mt_config.logical_axis_rules): trainer = MaxTextPeftTrainer(model, optimizer, tunix_config) - if mt_config.lora.lora_restore_path: - trainer = lora_utils.restore_lora_from_path(trainer, mt_config) + if mt_config.lora.lora_restore_path and trainer.train_steps == 0: + lora_utils.restore_lora_from_path(trainer.model, mt_config) trainer.with_training_hooks(training_hooks) trainer.with_data_hooks(data_hooks) trainer = use_maxtext_loss_function(trainer, mt_config) diff --git a/src/maxtext/trainers/pre_train/train.py b/src/maxtext/trainers/pre_train/train.py index e7f5ea84b6..8584ddcfa3 100644 --- a/src/maxtext/trainers/pre_train/train.py +++ b/src/maxtext/trainers/pre_train/train.py @@ -78,9 +78,7 @@ def get_first_step(model, state): if isinstance(model, nn.Module): return int(state.step) - if hasattr( - state, "inner_state" - ): # DiLoCoTrainState (NNX DiLoCo): step is the optimizer step var + if hasattr(state, "inner_state"): # DiLoCoTrainState (NNX DiLoCo): step is the optimizer step var return int(state.step.get_value()) return int(state.optimizer.step.get_value()) @@ -147,7 +145,7 @@ def loss_fn(model, config, data, dropout_rng, params, sparsity_state=None, is_tr encoder_images=data["images"] if config.use_multimodal else None, encoder_image_masks=data["image_masks"] if config.use_multimodal and "image_masks" in data else None, enable_dropout=config.enable_dropout if is_train else False, - rngs={"dropout": rng1, "params": aqt_rng}, + rngs={"dropout": rng1, "params": aqt_rng}, # pyrefly: ignore[bad-argument-type] mutable=mutable_collections, decoder_target_tokens=data["targets"], decoder_target_mask=data["targets_segmentation"], @@ -375,11 +373,11 @@ def train_step(model, config, state_mesh_shardings, params_shardings, state, dat if config.shard_optimizer_over_data: params = jax.tree.map( functools.partial(sharding.maybe_shard_with_name, shard_mode=config.shard_mode), - params, + params, # pyrefly: ignore[unbound-name] params_shardings, ) sparsity_enabled = config.weight_sparsity_n and config.weight_sparsity_m - pure_params = params["params"] if sparsity_enabled else params + pure_params = params["params"] if sparsity_enabled else params # pyrefly: ignore[unbound-name] batch_stats = params.get("batch_stats", {}) grad_func = jax.value_and_grad(loss_fn, argnums=4, has_aux=True) @@ -784,11 +782,7 @@ def train_loop(config, recorder, state=None): # the Zero-1 opt overlay doesn't apply through the diloco wrapper. params_shardings = state_mesh_shardings.params else: - params_shardings, state_mesh_shardings = ( - sharding.maybe_update_params_sharding_with_opt( - config, state_mesh_shardings - ) - ) + params_shardings, state_mesh_shardings = sharding.maybe_update_params_sharding_with_opt(config, state_mesh_shardings) p_train_step, p_eval_step = train_utils.jit_train_and_eval_step( config, @@ -828,9 +822,7 @@ def train_loop(config, recorder, state=None): if isinstance(model, nn.Module): setup_params = state.params elif config.enable_diloco: - setup_params = ( - state.params - ) # DiLoCoTrainState.params: the outer (global) params + setup_params = state.params # DiLoCoTrainState.params: the outer (global) params else: _, setup_params, _ = nnx.split(state.model, nnx.Param, ...) metric_logger_instance.write_setup_info_to_tensorboard(setup_params) @@ -935,6 +927,9 @@ def initialize(argv: Sequence[str]) -> tuple[pyconfig.HyperParameters, Any]: if config.use_vertex_tensorboard or os.environ.get("UPLOAD_DATA_TO_TENSORBOARD"): vertex_tensorboard_manager.configure_vertex_tensorboard(config) + if config.use_te_comm_gemm_overlap: + max_utils.bootstrap_transformer_engine_cgemm(config) + # Create the Goodput recorder recorder = create_goodput_recorder(config) diff --git a/src/maxtext/trainers/pre_train/train_compile.py b/src/maxtext/trainers/pre_train/train_compile.py index d472bb68d4..aa33ce73cb 100644 --- a/src/maxtext/trainers/pre_train/train_compile.py +++ b/src/maxtext/trainers/pre_train/train_compile.py @@ -187,7 +187,7 @@ def create_train_state_fn(): # Collect NNX activation shardings via an abstract forward pass (must run # after get_abstract_state, which only traces __init__). if config.debug_sharding and config.pure_nnx: - _collect_nnx_activation_shardings(_create_model_partial, config, topology_mesh) + _collect_nnx_activation_shardings(_create_model_partial, config, topology_mesh) # pyrefly: ignore[unbound-name] return ( shaped_train_args, diff --git a/src/maxtext/trainers/tokenizer/train_tokenizer.py b/src/maxtext/trainers/tokenizer/train_tokenizer.py index d3b4f35fce..58c63955fb 100644 --- a/src/maxtext/trainers/tokenizer/train_tokenizer.py +++ b/src/maxtext/trainers/tokenizer/train_tokenizer.py @@ -82,7 +82,7 @@ def build_grain_iterator(data_file_pattern: str, data_file_type: str, data_keys: if data_file_type == "parquet": dataset = grain.MapDataset.source(data_files) dataset = dataset.map(grain.experimental.ParquetIterDataset) - dataset = grain.experimental.InterleaveIterDataset(dataset, cycle_length=len(data_files)) + dataset = grain.experimental.InterleaveIterDataset(dataset, cycle_length=len(data_files)) # pyrefly: ignore[bad-argument-type] dataset = dataset.map(input_pipeline_utils.KeepFeatures(feature_names=list(data_keys))) return iter(dataset) elif data_file_type == "arrayrecord": @@ -94,7 +94,7 @@ def build_grain_iterator(data_file_pattern: str, data_file_type: str, data_keys: elif data_file_type == "tfrecord": dataset = grain.MapDataset.source(data_files) dataset = dataset.map(input_pipeline_utils.make_tfrecord_iter_dataset) - dataset = grain.experimental.InterleaveIterDataset(dataset, cycle_length=len(data_files)) + dataset = grain.experimental.InterleaveIterDataset(dataset, cycle_length=len(data_files)) # pyrefly: ignore[bad-argument-type] dataset = dataset.map(input_pipeline_utils.ParseFeatures(list(data_keys), tokenize=True)) dataset = dataset.map(input_pipeline_utils.NormalizeFeatures(list(data_keys), tokenize=True)) return iter(dataset) diff --git a/src/maxtext/utils/elastic_utils.py b/src/maxtext/utils/elastic_utils.py index 29e9cfb3fe..cef8d73054 100644 --- a/src/maxtext/utils/elastic_utils.py +++ b/src/maxtext/utils/elastic_utils.py @@ -16,6 +16,7 @@ import functools from collections import Counter +from types import SimpleNamespace import jax from maxtext.utils import gcs_utils @@ -226,3 +227,34 @@ def maybe_elastic_scale_up(config, checkpoint_manager): checkpoint_manager.wait_until_finished() max_logging.log("Checkpoint save completed. Interrupting") raise manager.ScaleUpSignalError() + + +def single_controller_mtc_init_kwargs(raw_keys): + """Returns topology kwargs for single-controller MTC initialization.""" + kwargs = { + "data_parallelism": raw_keys["mtc_data_parallelism"], + "num_slices": raw_keys["num_slices"], + } + if not raw_keys.get("elastic_enabled", False): + return kwargs + + config = SimpleNamespace(**raw_keys) + if not should_use_elastic(config): + return kwargs + + active_devices = tuple(live_devices(config)) + active_slice_indices = {getattr(device, "slice_index", 0) for device in active_devices if device is not None} + if not active_devices or not active_slice_indices: + raise ValueError("Elastic single-controller MTC initialization found no active devices.") + + kwargs["devices"] = active_devices + kwargs["num_slices"] = len(active_slice_indices) + if not kwargs["data_parallelism"]: + kwargs["data_parallelism"] = kwargs["num_slices"] + max_logging.log( + "Using active elastic devices for single-controller MTC initialization: " + f"active_num_slices={kwargs['num_slices']}, " + f"active_device_count={len(active_devices)}, " + f"configured_num_slices={raw_keys['num_slices']}." + ) + return kwargs diff --git a/src/maxtext/utils/generate_param_only_checkpoint.py b/src/maxtext/utils/generate_param_only_checkpoint.py index 8a2488abe1..a44a7819f2 100644 --- a/src/maxtext/utils/generate_param_only_checkpoint.py +++ b/src/maxtext/utils/generate_param_only_checkpoint.py @@ -49,6 +49,23 @@ from maxtext.utils import model_creation_utils from maxtext.utils import train_utils +_SCAN_AXIS_NAMES = ("layers_per_stage", "layers", "circular_repeats") + + +def get_scan_axis(config, sharding_or_spec, x=None): + """Find the scan axis from sharding spec or use default fallback if valid.""" + if hasattr(sharding_or_spec, "spec"): + spec = sharding_or_spec.spec + else: + spec = sharding_or_spec + + for name in _SCAN_AXIS_NAMES: + if name in spec: + return spec.index(name) + if x is not None and hasattr(x, "ndim") and x.ndim > config.param_scan_axis: + return config.param_scan_axis + return None + def _possibly_unroll_params(config, training_state, training_state_annotations, mesh): """Unroll scanned input layers when force_unroll is set.""" @@ -106,8 +123,17 @@ def _possibly_unroll_params_nnx(config, state, state_mesh_shardings, mesh): and removes the original collection. Mirrors the same operation on `state_mesh_shardings` so downstream sharding stays correct. """ - decoder_state = state.model.decoder - decoder_shardings = state_mesh_shardings.model.decoder + param_only_state, other_state = nnx.split_state(state, nnx.Param, ...) + param_paths = set(param_only_state.flat_state().paths) + + param_only_shardings_items = [(path, val) for path, val in state_mesh_shardings.flat_state() if path in param_paths] + other_shardings_items = [(path, val) for path, val in state_mesh_shardings.flat_state() if path not in param_paths] + + param_only_shardings = nnx.State.from_flat_path(param_only_shardings_items) + other_shardings = nnx.State.from_flat_path(other_shardings_items) + + decoder_state = param_only_state.model.decoder + decoder_shardings = param_only_shardings.model.decoder def unroll_layer_group(num_layers, layer_name="layers"): layers = decoder_state.get(layer_name, None) @@ -115,19 +141,42 @@ def unroll_layer_group(num_layers, layer_name="layers"): if layers is None or layers_shardings is None: raise ValueError(f"Missing {layer_name} in NNX state.model.decoder or state_mesh_shardings.") - def drop_scan_axis(named_sharding): + def drop_scan_axis(named_sharding, x): ps = named_sharding.spec - return jax.sharding.PartitionSpec(*(ps[0 : config.param_scan_axis] + ps[config.param_scan_axis + 1 :])) + val = x[...] if isinstance(x, nnx.Variable) else x + scan_axis = get_scan_axis(config, named_sharding, val) + if scan_axis is not None: + return jax.sharding.PartitionSpec(*(ps[0:scan_axis] + ps[scan_axis + 1 :])) + return ps new_layer_pspec = jax.tree_util.tree_map( - drop_scan_axis, layers_shardings, is_leaf=lambda x: isinstance(x, jax.sharding.NamedSharding) + drop_scan_axis, + layers_shardings, + layers, + is_leaf=lambda x: isinstance(x, (jax.sharding.NamedSharding, nnx.Variable)), ) new_layer_sharding = jax.tree_util.tree_map(lambda ps: jax.sharding.NamedSharding(mesh, ps), new_layer_pspec) for i in range(num_layers): + # pylint: disable=cell-var-from-loop def slice_ith(input_layers): - return jax.tree_util.tree_map(lambda x: jnp.take(x, i, axis=config.param_scan_axis), input_layers) + def _slice_leaf(x, sharding): + val = x[...] if isinstance(x, nnx.Variable) else x + scan_axis = get_scan_axis(config, sharding, val) + if scan_axis is not None: + sliced_val = jnp.take(val, i, axis=scan_axis) + if isinstance(x, nnx.Variable): + return type(x)(sliced_val) + return sliced_val + return x + + return jax.tree_util.tree_map( + _slice_leaf, + input_layers, + layers_shardings, + is_leaf=lambda x: isinstance(x, (jax.sharding.NamedSharding, nnx.Variable)), + ) # pylint: disable=not-callable new_layer = jax.jit(slice_ith, out_shardings=new_layer_sharding)(layers) @@ -145,6 +194,13 @@ def slice_ith(input_layers): else: unroll_layer_group(config.num_decoder_layers, layer_name="layers") + # Merge modified parameter state and other variables back in-place + merged_state = nnx.merge_state(param_only_state, other_state) + state.update(merged_state) + + merged_shardings = nnx.merge_state(param_only_shardings, other_shardings) + state_mesh_shardings.update(merged_shardings) + def _read_train_checkpoint(config, checkpoint_manager, mesh): """Read training checkpoint at path defined by load_full_state_path.""" @@ -167,13 +223,14 @@ def init_state_fn(): tx = optimizers.get_optimizer(config, learning_rate_schedule) init_state_fn = functools.partial(maxtext_utils.init_initial_state, model, tx, config, True, rng) - state, state_mesh_notations, state_mesh_shardings, _ = maxtext_utils.setup_training_state( + state, state_mesh_notations, state_mesh_shardings, _, _ = maxtext_utils.setup_training_state( None, config, mesh, checkpoint_manager, init_state_fn ) if config.pure_nnx: # On NNX, state is a flat nnx.State; params live under state.model and the # legacy notations are unused (callers receive shardings directly). - num_params = max_utils.calculate_num_params_from_pytree(state.model) + params, _ = nnx.split_state(state.model, nnx.Param, ...) + num_params = max_utils.calculate_num_params_from_pytree(params) max_logging.log(f"In input checkpoint Number of model params={num_params/1e9:.3f} billion") return state, state_mesh_shardings num_params = max_utils.calculate_num_params_from_pytree(state.params) @@ -283,16 +340,36 @@ def unroll_layer_group(num_layers, layer_name="layers"): if layers is None or layers_annotations is None: return # No LoRA on this layer group; nothing to unroll. - def new_pspec(x): - return jax.sharding.PartitionSpec(*(x[0 : config.param_scan_axis] + x[config.param_scan_axis + 1 :])) - - new_layer_annotation = jax.tree_util.tree_map(new_pspec, layers_annotations) + def drop_scan_axis(spec, x): + scan_axis = get_scan_axis(config, spec, x) + if scan_axis is not None: + return jax.sharding.PartitionSpec(*(spec[0:scan_axis] + spec[scan_axis + 1 :])) + return spec + + new_layer_annotation = jax.tree_util.tree_map( + drop_scan_axis, + layers_annotations, + layers, + is_leaf=lambda x: isinstance(x, jax.sharding.PartitionSpec), + ) new_layer_sharding = jax.tree_util.tree_map(lambda x: jax.sharding.NamedSharding(mesh, x), new_layer_annotation) for i in range(num_layers): + # pylint: disable=cell-var-from-loop def slice_ith(input_layers): - return jax.tree_util.tree_map(lambda x: jnp.take(x, i, axis=config.param_scan_axis), input_layers) + def _slice_leaf(x, spec): + scan_axis = get_scan_axis(config, spec, x) + if scan_axis is not None: + return jnp.take(x, i, axis=scan_axis) + return x + + return jax.tree_util.tree_map( + _slice_leaf, + input_layers, + layers_annotations, + is_leaf=lambda x: isinstance(x, jax.sharding.PartitionSpec), + ) # pylint: disable=not-callable new_layer = jax.jit(slice_ith, out_shardings=new_layer_sharding)(layers) @@ -400,13 +477,9 @@ def generate_decode_checkpoint(config): training_state, training_state_annotations = _read_train_checkpoint(config, checkpoint_manager, mesh) if config.pure_nnx: # NNX state is a flat nnx.State; opt_state lives under the optimizer sub-state. - assert ( - training_state.optimizer.opt_state - ), "missing opt_state in training checkpoint" + assert training_state.optimizer.opt_state, "missing opt_state in training checkpoint" else: - assert ( - training_state.opt_state != {} - ), "missing opt_state in training checkpoint" + assert training_state.opt_state != {}, "missing opt_state in training checkpoint" _possibly_unroll_params(config, training_state, training_state_annotations, mesh) diff --git a/src/maxtext/utils/gradient_accumulation.py b/src/maxtext/utils/gradient_accumulation.py index 5c76588b64..35fdf65503 100644 --- a/src/maxtext/utils/gradient_accumulation.py +++ b/src/maxtext/utils/gradient_accumulation.py @@ -16,6 +16,7 @@ import jax import jax.numpy as jnp +from jax.experimental.xla_metadata import set_xla_metadata from jax.sharding import NamedSharding from flax import nnx @@ -108,11 +109,16 @@ def convert_to_bf16(param): def accumulate_gradient(acc_grad_and_loss, data): ga_params = acc_grad_and_loss["ga_params"] + # Scope double-buffering to model scans, not the enclosing gradient accumulation scan. + # This metadata only affects XLA:GPU; see https://github.com/openxla/xla/pull/43468. + # Enable it with JAX_OPTIMIZATION_LEVEL=O1 and + # XLA_FLAGS="--xla_gpu_enable_while_loop_unrolling=WHILE_LOOP_UNROLLING_MANUAL_UNROLL". if is_nnx: # Reconstruct the model using the fixed parameters (ga_params) # and the advancing non-parameter state (RNGs) from the carry. local_model = nnx.merge(graphdef, ga_params, acc_grad_and_loss["rest_state"], copy=True) - (_, aux), cur_batch_gradient = grad_func(local_model, config, data, None, None, is_train=True) + with set_xla_metadata(_xla_loop_unroll_strategy="double-buffer"): + (_, aux), cur_batch_gradient = grad_func(local_model, config, data, None, None, is_train=True) _, _, next_rest_state = nnx.split(local_model, nnx.Param, ...) acc_grad_and_loss["rest_state"] = next_rest_state else: @@ -121,7 +127,8 @@ def accumulate_gradient(acc_grad_and_loss, data): if dropout_rng is not None else None ) - (_, aux), cur_batch_gradient = grad_func(model, config, data, rng, ga_params, is_train=True) + with set_xla_metadata(_xla_loop_unroll_strategy="double-buffer"): + (_, aux), cur_batch_gradient = grad_func(model, config, data, rng, ga_params, is_train=True) acc_grad_and_loss["loss"] += aux["xent_sum"] + aux.get("dpo_loss", 0.0) acc_grad_and_loss["moe_lb_loss"] += aux["moe_lb_loss"] acc_grad_and_loss["indexer_loss"] += aux["indexer_loss"] @@ -150,7 +157,7 @@ def reshape_to_microbatch_accumulations(batch_arr): "ga_params": ga_params, } if is_nnx: - init_grad_and_loss["rest_state"] = rest + init_grad_and_loss["rest_state"] = rest # pyrefly: ignore[unbound-name] grad_and_loss, aux = jax.lax.scan( accumulate_gradient, init_grad_and_loss, data, length=config.gradient_accumulation_steps diff --git a/src/maxtext/utils/layerwise_quantization.py b/src/maxtext/utils/layerwise_quantization.py index 3279b59a6f..f8df26967a 100644 --- a/src/maxtext/utils/layerwise_quantization.py +++ b/src/maxtext/utils/layerwise_quantization.py @@ -339,7 +339,12 @@ def _load_and_quantize_nnx(self) -> None: # Convert-mode state has both `kernel` (full precision) and `AqtDotGeneral_0.qrhs.frozen` # at every quantized DenseGeneral; the serve-mode reader expects only the latter. - convert_state = nnx.state(convert_model).to_pure_dict() + # We filter out transient runtime state (RngState, Cache, Intermediate, BatchStat) so only persistent weights are saved. + convert_state = ( + nnx.state(convert_model) + .filter(lambda path, var: not isinstance(var, (nnx.RngState, nnx.Cache, nnx.Intermediate, nnx.BatchStat))) + .to_pure_dict() + ) serve_state = self._strip_kernels_at_quantized_paths(convert_state) if config.save_quantized_params_path: diff --git a/src/maxtext/utils/lora_utils.py b/src/maxtext/utils/lora_utils.py index 6b4410f209..1aa443f093 100644 --- a/src/maxtext/utils/lora_utils.py +++ b/src/maxtext/utils/lora_utils.py @@ -19,7 +19,7 @@ import json import os import re -from typing import Any, Optional +from typing import Optional from flax import nnx, linen as nn from flax.linen import partitioning as nn_partitioning @@ -465,10 +465,9 @@ def _build_lora_provider(mt_config: pyconfig.HyperParameters) -> qwix.LoraProvid return qwix.LoraProvider(**lora_kwargs) -def _prepare_dummy_inputs() -> tuple[jnp.ndarray, jnp.ndarray]: +def _prepare_dummy_inputs(dummy_bs: int = 1) -> tuple[jnp.ndarray, jnp.ndarray]: """Builds dummy decoder inputs used to materialize LoRA parameters.""" # Keep LoRA warmup as small as possible to minimize compile/memory overhead. - dummy_bs = 1 seq_len = 1 decoder_input_tokens = jnp.zeros((dummy_bs, seq_len), dtype=jnp.int32) decoder_positions = jnp.zeros((dummy_bs, seq_len), dtype=jnp.int32) @@ -483,30 +482,50 @@ def is_lora_enabled(model: nnx.Module) -> bool: return False -def _verify_lora_parameters(lora_model: nnx.Module, mt_config: pyconfig.HyperParameters): +def _verify_lora_parameters(lora_model: nnx.Module, mt_config: pyconfig.HyperParameters) -> None: """Validates that LoRA is active or that target modules were matched.""" if is_lora_enabled(lora_model): + wrapped_modules = set() + for path, value in nnx.iter_graph(lora_model): + if isinstance(value, nnx.LoRAParam): + if len(path) > 1: + parent_path = "/".join(str(p) for p in path[:-1]) + wrapped_modules.add(parent_path) + + if wrapped_modules: + wrapped_modules = sorted(list(wrapped_modules)) + max_logging.log( + f"LoRA configured: module_path='{_get_lora_module_path(mt_config)}' successfully matched " + f"{len(wrapped_modules)} target submodules." + ) + preview_limit = 20 + preview_modules = wrapped_modules[:preview_limit] + max_logging.log(f"Sample matched submodules ({len(preview_modules)} of {len(wrapped_modules)}): {preview_modules}") + else: + max_logging.log("LoRA is enabled. (Detailed submodules match report skipped due to mock model or empty state)") return lora_module_path = _get_lora_module_path(mt_config) compiled_module_path = re.compile(lora_module_path) - matched_module_paths = [] - sample_module_paths = [] + matched_module_paths = [] for path, _ in nnx.iter_modules(lora_model): module_path = "/".join(str(p) for p in path) - if len(sample_module_paths) < 100: - sample_module_paths.append(module_path) - if compiled_module_path.search(module_path): + if module_path and compiled_module_path.search(module_path): matched_module_paths.append(module_path) if not matched_module_paths: - max_logging.log( - f"LoRA module_path='{lora_module_path}' did not match any weights. " f"Sample module paths: {sample_module_paths}" - ) + max_logging.log(f"Error: LoRA module_path='{lora_module_path}' did not match any weights.") raise ValueError("LoRA enabled but no LoRA parameters found in decoder/model state.") + # Simplify matched paths by replacing numeric layer indices with "*" to avoid redundant output + simplified_matches = sorted( + {"/".join("*" if p.isdigit() else p for p in path.split("/")) for path in matched_module_paths} + ) + max_logging.log(f"LoRA target verification: successfully matched {len(matched_module_paths)} modules.") + max_logging.log(f"Matched submodule patterns: {simplified_matches}") + raise ValueError( "LoRA module path matched target modules, but nnx.LoRAParam is still " "missing. For Tunix PeftTrainer, LoRA params must be materialized before " @@ -515,6 +534,45 @@ def _verify_lora_parameters(lora_model: nnx.Module, mt_config: pyconfig.HyperPar ) +def sync_lora_metadata(config: pyconfig.HyperParameters) -> None: + """Syncs LoRA parameters (rank, alpha) from the checkpoint sidecar metadata if present. + + If configuration values are set to non-default values (i.e. rank > 0 or alpha > 0.0) + and differ from the checkpoint metadata values, we raise a ValueError to fail the run. + If they are at default values, we sync them from the checkpoint. + """ + lora_restore_path = config.lora.lora_restore_path + if not lora_restore_path: + return + + custom_metadata = checkpointing.load_checkpoint_metadata(lora_restore_path) + lora_meta = custom_metadata.get("lora") + + if lora_meta: + meta_rank = lora_meta.get("lora_rank", config.lora.lora_rank) + meta_alpha = lora_meta.get("lora_alpha", config.lora.lora_alpha) + + # Check lora_rank + if config.lora.lora_rank not in (0, meta_rank): + raise ValueError( + f"Configured lora_rank ({config.lora.lora_rank}) does not match " + f"checkpoint metadata lora_rank ({meta_rank}) at {lora_restore_path}." + ) + # Check lora_alpha + if config.lora.lora_alpha not in (0.0, meta_alpha): + raise ValueError( + f"Configured lora_alpha ({config.lora.lora_alpha}) does not match " + f"checkpoint metadata lora_alpha ({meta_alpha}) at {lora_restore_path}." + ) + + config.lora.lora_rank = meta_rank + config.lora.lora_alpha = meta_alpha + max_logging.log( + f"Synced LoRA parameters from Orbax metadata at {lora_restore_path}: " + f"rank={config.lora.lora_rank}, alpha={config.lora.lora_alpha}" + ) + + def apply_lora_to_model( model: nnx.Module, mesh: Optional[jax.sharding.Mesh], @@ -533,8 +591,12 @@ def apply_lora_to_model( lora_provider = _build_lora_provider(mt_config) + dp_size = 1 + if mesh is not None and "data" in mesh.shape: + dp_size = mesh.shape["data"] + model_rngs = getattr(model.decoder, "rngs", None) - decoder_input_tokens, decoder_positions = _prepare_dummy_inputs() + decoder_input_tokens, decoder_positions = _prepare_dummy_inputs(dummy_bs=dp_size) lora_model = qwix.apply_lora_to_model( model, @@ -582,18 +644,27 @@ def _safe_reshard(var, sharding_spec): return lora_model -def restore_lora_from_path(trainer: Any, mt_config: pyconfig.HyperParameters) -> Any: - """Restores LoRA parameter weights from an external Orbax checkpoint for a fresh run.""" - lora_restore_path = mt_config.lora.lora_restore_path +def restore_lora_from_path(model: nnx.Module, mt_config: pyconfig.HyperParameters) -> nnx.Module: + """Restores LoRA parameter weights from an external Orbax checkpoint. - train_steps = getattr(trainer, "train_steps", 0) - if train_steps > 0: - max_logging.log( - f"PeftTrainer restored current run at step {train_steps}; " f"ignoring lora_restore_path '{lora_restore_path}'." - ) - return trainer + This function performs the restore in-place on the model's parameters and + returns the model with the restored weights applied. - if not is_lora_enabled(trainer.model): + Args: + model: The JAX/Flax NNX model (nnx.Module). + mt_config: The HyperParameters config containing the lora configuration. + + Returns: + The model with the restored LoRA weights applied in-place. + + Raises: + ValueError: If LoRA is not enabled on the model, but a restore path is set. + """ + lora_restore_path = mt_config.lora.lora_restore_path + if not lora_restore_path: + return model + + if not is_lora_enabled(model): lora_module_path = _get_lora_module_path(mt_config) if not mt_config.lora.enable_lora: raise ValueError( @@ -601,7 +672,9 @@ def restore_lora_from_path(trainer: Any, mt_config: pyconfig.HyperParameters) -> f"Set lora.enable_lora=True and verify lora_module_path ('{lora_module_path}') matches model modules." ) - abstract_lora_params = nnx.state(trainer.model, nnx.LoRAParam) + sync_lora_metadata(mt_config) + + abstract_lora_params = nnx.state(model, nnx.LoRAParam) target_for_restore = jax.tree.map( lambda v: {"value": v.value}, @@ -657,9 +730,9 @@ def _map_to_state(path, variable): is_leaf=lambda n: isinstance(n, nnx.Variable), ) - nnx.update(trainer.model, abstract_lora_params) + nnx.update(model, abstract_lora_params) max_logging.log(f"LoRA restore complete from '{lora_restore_path}'.") - return trainer + return model # NNX-shaped LoRA helpers. diff --git a/src/maxtext/utils/max_utils.py b/src/maxtext/utils/max_utils.py index 4ee35a3844..98db971ef2 100644 --- a/src/maxtext/utils/max_utils.py +++ b/src/maxtext/utils/max_utils.py @@ -248,14 +248,14 @@ def maybe_initialize_jax_distributed_system(raw_keys): max_logging.log("Skipping jax distributed system since its not needed for single controller.") if raw_keys["enable_multi_tier_checkpointing"]: max_logging.log("Initializing multi-tier checkpointing for single controller...") + mtc_init_kwargs = elastic_utils.single_controller_mtc_init_kwargs(raw_keys) initialize_multi_tier_checkpointing( local_checkpoint_directory=raw_keys["local_checkpoint_directory"], backup_interval_minutes=raw_keys["multi_tier_checkpointing_backup_interval_minutes"], run_name=raw_keys["run_name"], jax_initialization_timeout_seconds=raw_keys["jax_distributed_initialization_timeout"], - data_parallelism=raw_keys["mtc_data_parallelism"], - num_slices=raw_keys["num_slices"], use_colocated_python=True, + **mtc_init_kwargs, ) return if jax.distributed.is_initialized(): @@ -336,8 +336,8 @@ def initialize_jax_for_gpu(raw_keys): jax.distributed.initialize( coordinator_address=f"{coordinator_ip}:{coordinator_port}", - num_processes=int(os.getenv("NNODES")), - process_id=int(os.getenv("NODE_RANK")), + num_processes=int(os.getenv("NNODES")), # pyrefly: ignore[bad-argument-type] + process_id=int(os.getenv("NODE_RANK")), # pyrefly: ignore[bad-argument-type] initialization_timeout=raw_keys["jax_distributed_initialization_timeout"], local_device_ids=devices, ) @@ -349,16 +349,16 @@ def initialize_jax_for_cpu(raw_keys): coordinator_ip_address = get_coordinator_ip_address() coordinator_address = coordinator_ip_address + ":1234" # JAX coordinator port used in XPK # Env variables to be set in XPK or otherwise - job_index = int(os.environ.get("JOB_INDEX")) - job_completion_index = int(os.environ.get("JOB_COMPLETION_INDEX")) - processes_in_job = int(os.environ.get("PROCESSES_IN_JOB")) + job_index = int(os.environ.get("JOB_INDEX")) # pyrefly: ignore[bad-argument-type] + job_completion_index = int(os.environ.get("JOB_COMPLETION_INDEX")) # pyrefly: ignore[bad-argument-type] + processes_in_job = int(os.environ.get("PROCESSES_IN_JOB")) # pyrefly: ignore[bad-argument-type] pid = job_index * processes_in_job + job_completion_index max_logging.log(f" Jax process id is {pid} ") # Explicit initialize is needed only for CPUs jax.distributed.initialize( coordinator_address=coordinator_address, process_id=pid, - num_processes=int(os.environ.get("JAX_PROCESS_COUNT")), + num_processes=int(os.environ.get("JAX_PROCESS_COUNT")), # pyrefly: ignore[bad-argument-type] initialization_timeout=raw_keys["jax_distributed_initialization_timeout"], ) @@ -444,7 +444,7 @@ def get_coordinator_ip_address(): max_coordinator_lookups = 50 while not coordinator_found and lookup_attempt <= max_coordinator_lookups: try: - coordinator_ip_address = socket.gethostbyname(coordinator_address) + coordinator_ip_address = socket.gethostbyname(coordinator_address) # pyrefly: ignore[bad-argument-type] coordinator_found = True except socket.gaierror: max_logging.log( @@ -676,7 +676,7 @@ def _cross_entropy_with_logits_fwd(logits: jnp.ndarray, targets: jnp.ndarray, z_ log_z = jnp.squeeze(jnp.log(sum_exp) + max_logit, axis=-1) total_z_loss = z_loss * jax.lax.square(log_z) loss += total_z_loss - return (loss, total_z_loss), ( + return (loss, total_z_loss), ( # pyrefly: ignore[bad-return] logits, targets, z_loss, @@ -698,11 +698,11 @@ def _cross_entropy_with_logits_bwd( g: tuple[jnp.ndarray, jnp.ndarray], ) -> tuple[jnp.ndarray, None, None]: """Backward-mode of `cross_entropy_with_logits`.""" - g = g[0] # Ignore z_loss component as that is only used for logging. + g = g[0] # Ignore z_loss component as that is only used for logging. # pyrefly: ignore[bad-assignment] logits, targets, z_loss, exp_shifted, sum_exp, log_z = res # z-loss term adds the (2 * z_loss * log_z) factor. deriv = jnp.expand_dims(1 + 2 * z_loss * log_z, -1) * exp_shifted / sum_exp - targets - g_logits = jnp.expand_dims(g, axis=-1) * deriv + g_logits = jnp.expand_dims(g, axis=-1) * deriv # pyrefly: ignore[bad-argument-type] return ( jnp.asarray(g_logits, logits.dtype), @@ -1185,6 +1185,22 @@ def print_non_trivial_mesh_axis(mesh): print(f"{mesh_axis}: {axis_size}", flush=True) +def bootstrap_transformer_engine_cgemm(config): + """Potentially initialize NCCL communicators for Collective GEMM operations if + the environment is distributed and has the appropriate config.""" + import transformer_engine.jax.cpp_extensions as tex # pylint: disable=import-outside-toplevel # pytype: disable=import-error + + tsp_size = config.ici_tensor_sequence_parallelism * config.dcn_tensor_sequence_parallelism + + # Setup NCCL buffers for GPU Collective GEMM operations + tex.collective_gemm_bootstrap( + jax.device_count(), + jax.local_device_count(), + jax.process_index(), + tsp_size, + ) + + @contextmanager def maybe_get_transformer_engine_context(config): """Runs a transformer engine context engine manager for GPUs only.""" @@ -1212,9 +1228,9 @@ def transformer_engine_context(): mesh_resource = MeshResource( # pytype: disable=wrong-arg-types dp_resource="data", tp_resource="tensor", - # tpsp_resource = "tensor_sequence", #TODO(Phuong): add this back when upstreaming CGEMM + tpsp_resource="tensor_sequence", fsdp_resource="fsdp", - pp_resource=None, + pp_resource=None, # pyrefly: ignore[bad-argument-type] cp_resource="context", ) with global_shard_guard(mesh_resource): diff --git a/src/maxtext/utils/maxtext_utils.py b/src/maxtext/utils/maxtext_utils.py index 3451ec824d..fe7fa5a426 100644 --- a/src/maxtext/utils/maxtext_utils.py +++ b/src/maxtext/utils/maxtext_utils.py @@ -95,7 +95,7 @@ def get_functional_train_with_signature( ): """Get the shardings (both state and data) for `train_step`.""" functional_train = functools.partial(train_step, model, config, state_mesh_shardings, params_shardings) - functional_train.__name__ = "train_step" + functional_train.__name__ = "train_step" # pyrefly: ignore[missing-attribute] if config.pure_nnx: in_shardings = (state_mesh_shardings, data_sharding) # State, batch else: @@ -109,7 +109,7 @@ def get_functional_train_with_signature( def get_functional_eval_with_signature(eval_step, data_sharding, state_mesh_shardings, model, config): """Get the shardings (both state and data) for `eval_step`.""" functional_eval = functools.partial(eval_step, model, config) - functional_eval.__name__ = "eval_step" + functional_eval.__name__ = "eval_step" # pyrefly: ignore[missing-attribute] if config.pure_nnx: in_shardings = (state_mesh_shardings, data_sharding) # State, batch (NNX: no rng) else: @@ -1392,7 +1392,9 @@ def get_abstract_param(model, config): {"params": key, "dropout": key, "aqt": key}, np.ones(input_shape, dtype=jnp.int32), np.ones(input_shape, dtype=jnp.int32), - encoder_images=np.ones(image_shape, dtype=jnp.int32) if config.use_multimodal else None, + encoder_images=np.ones(image_shape, dtype=jnp.int32) + if config.use_multimodal + else None, # pyrefly: ignore[no-matching-overload] encoder_audios=np.ones(audio_shape, dtype=jnp.float32) if config.use_audio else None, ) return abstract_vars @@ -1413,7 +1415,7 @@ def setup_decode_state(config, mesh, checkpoint_manager, init_state_fn): if not config.load_parameters_path: # generate random params max_logging.log("No decode checkpoint specified - generating random weights.") - state, state_mesh_annotations, _, _ = setup_initial_state( + state, state_mesh_annotations, _, _, _ = setup_initial_state( None, config, mesh, checkpoint_manager, init_state_fn, False ) else: @@ -1468,6 +1470,9 @@ def setup_initial_state( Returns: train_state: the initialized train state. For NNX, this is a TrainStateNNX instance state_mesh_annotations: the mesh annotations for the train state + state_mesh_shardings: the mesh shardings for the train state + data_iterator: the updated data iterator + was_restored: True if state or params were restored from checkpoint, False if freshly initialized """ unboxed_abstract_state, state_mesh_annotations, state_mesh_shardings = get_abstract_state( @@ -1493,6 +1498,8 @@ def setup_initial_state( expansion_factor_real_data=config.expansion_factor_real_data, maxtext_config=config, ) + # Partial or fully restored + was_restored = bool(restored is not None or raw_params is not None) if restored: if isinstance( @@ -1548,7 +1555,7 @@ def _merge_params(p_raw, p_init): state = state.replace(params=raw_params) if not config.pure_nnx: state = max_utils.unbox_logicallypartioned(state) - return state, state_mesh_annotations, state_mesh_shardings, data_iterator + return state, state_mesh_annotations, state_mesh_shardings, data_iterator, was_restored def get_logical_annotations(config, mesh, init_state_fn): @@ -2022,7 +2029,7 @@ def to_abstract(x): # Convert all input arguments recursively to purely local abstract ShapeDtypeStruct objects # to completely bypass remote Array objects and proxy tracing overhead. abstract_inputs = jax.tree.map(to_abstract, train_step_inputs) - p_train_jaxpr = jax.make_jaxpr(unwrapped_step)(*abstract_inputs) + p_train_jaxpr = jax.make_jaxpr(unwrapped_step)(*abstract_inputs) # pyrefly: ignore[no-matching-overload] local_filename = "train_step.jaxpr" local_path = os.path.join(config.dump_jaxpr_local_dir, local_filename) @@ -2067,3 +2074,43 @@ def get_mesh_from_config( axis_types = tuple([AxisType.Auto] * len(config.mesh_axes)) return Mesh(devices_array, config.mesh_axes, axis_types=axis_types) + + +def prepare_kv_caches_for_scan(kv_caches, scan_length, block_len, stack=False): + """Groups the flat list of KV caches into block-sized tuples and optionally stacks them.""" + if kv_caches is None: + return None + + if not isinstance(kv_caches, list): + raise TypeError(f"kv_caches must be a list, got {type(kv_caches)}") + + grouped = [tuple(kv_caches[i * block_len : (i + 1) * block_len]) for i in range(scan_length)] + + if stack: + # Stack the list of tuples into a tuple of stacked PyTrees along a new leading axis + return jax.tree_util.tree_map(lambda *args: jnp.stack(args, axis=0), *grouped) + else: + return grouped + + +def update_kv_caches_after_scan(kv_caches, returned_kv_cache, scan_length, block_len, stacked=False): + """Updates the original flat list of KV caches from the scanned outputs.""" + if kv_caches is None or returned_kv_cache is None: + return + + if not isinstance(kv_caches, list): + raise TypeError(f"kv_caches must be a list, got {type(kv_caches)}") + + if stacked: + # Unstack the returned tuple of stacked PyTrees back into a list of tuples + unstacked = [jax.tree_util.tree_map(lambda x, i=i: x[i], returned_kv_cache) for i in range(scan_length)] + for i in range(scan_length): + start_idx = i * block_len + for offset, updated_item in enumerate(unstacked[i]): + kv_caches[start_idx + offset] = updated_item + else: + # returned_kv_cache is already a list of tuples + for i in range(scan_length): + start_idx = i * block_len + for offset, updated_item in enumerate(returned_kv_cache[i]): + kv_caches[start_idx + offset] = updated_item diff --git a/src/maxtext/utils/maxtext_utils_nnx.py b/src/maxtext/utils/maxtext_utils_nnx.py index 20494dbc71..7b265ab3bd 100644 --- a/src/maxtext/utils/maxtext_utils_nnx.py +++ b/src/maxtext/utils/maxtext_utils_nnx.py @@ -164,7 +164,7 @@ def create_nnx_sharded_model( named_sharding = nnx_extract_named_sharding(abstract_state) if mesh is None: - mesh = abstract_model.mesh + mesh = abstract_model.mesh # pyrefly: ignore[missing-attribute] # JIT a function that creates the model state with proper sharding from the start. # By providing out_shardings, we instruct JAX to produce sharded output directly, diff --git a/src/maxtext/utils/model_creation_utils.py b/src/maxtext/utils/model_creation_utils.py index c0248783eb..a8d62058e2 100644 --- a/src/maxtext/utils/model_creation_utils.py +++ b/src/maxtext/utils/model_creation_utils.py @@ -41,10 +41,11 @@ from flax import nnx from flax.core.meta import Partitioned import flax.linen as nn +from huggingface_hub import get_token import jax import jax.numpy as jnp from jax.sharding import Mesh -from maxtext.common.checkpointing import handle_checkpoint_mismatch +from maxtext.common import checkpointing from maxtext.common.common_types import MODEL_MODE_AUTOREGRESSIVE, MODEL_MODE_TRAIN from maxtext.configs import pyconfig from maxtext.integration.tunix.tunix_adapter import TunixMaxTextAdapter @@ -790,6 +791,45 @@ def create_models_and_meshes(trainer_config, sampler_config, trainer_devices, sa return reference_model, reference_mesh, actor_model, actor_mesh, rollout_mesh +def verify_and_sync_scan_layers(config): + """Verify and sync scan_layers based on checkpoint metadata.""" + if not config.load_parameters_path: + return config + + custom_metadata = checkpointing.load_checkpoint_metadata(config.load_parameters_path) + saved_scan_layers = custom_metadata.get("scan_layers") + if not isinstance(saved_scan_layers, bool): + return config + + if saved_scan_layers == config.scan_layers: + return config + + # Extract the Pydantic config (or use direct config if already a Pydantic model) + pydantic_config = getattr(config, "_pydantic_config", config) + model_fields_set = getattr(pydantic_config, "model_fields_set", None) + + # If model metadata tracking isn't supported, fall back to matching check (True) + is_explicit = "scan_layers" in model_fields_set if model_fields_set is not None else True + + if is_explicit: + if saved_scan_layers != config.scan_layers: + raise ValueError( + f"Configuration mismatch: Your run specifies scan_layers={config.scan_layers}, " + f"but the checkpoint was saved with scan_layers={saved_scan_layers}." + ) + else: + max_logging.log(f"Setting scan_layers={saved_scan_layers} loaded from checkpoint metadata.") + new_pydantic_config = pydantic_config.model_copy(update={"scan_layers": saved_scan_layers}) + # Wrap back in HyperParameters if the original config was wrapped + if getattr(config, "_pydantic_config", None) is not None: + config = pyconfig.HyperParameters(new_pydantic_config) + else: + config = new_pydantic_config + + return config + + +# pylint: disable=too-many-positional-arguments def from_pretrained( config, mesh=None, @@ -811,8 +851,12 @@ def from_pretrained( if config.convert_checkpoint_if_possible and not config.load_parameters_path: if not (epath.Path(config.base_output_directory) / "0" / "items").exists(): # Try to convert checkpoint on the fly - if not config.hf_access_token: - raise ValueError("hf_access_token must be provided when not providing a pre-existing checkpoint") + hf_access_token = config.hf_access_token or get_token() + if not hf_access_token: + raise ValueError( + "hf_access_token must be provided (or authenticate via" + " huggingface-cli) when not providing a pre-existing checkpoint" + ) # Only process 0 performs the conversion; other processes wait at the barrier below. # Otherwise every host would race to download from HF and concurrently write the same @@ -830,8 +874,8 @@ def from_pretrained( conversion_env = os.environ.copy() conversion_env["JAX_PLATFORMS"] = "cpu" # conversion_env["XLA_FLAGS"] = f"--xla_force_host_platform_device_count={simulated_cpu_devices_count}" - if config.hf_access_token: - conversion_env["HF_TOKEN"] = config.hf_access_token + if hf_access_token: + conversion_env["HF_TOKEN"] = hf_access_token to_maxtext_cmd = [ sys.executable, @@ -865,28 +909,24 @@ def from_pretrained( ) config = pyconfig.HyperParameters(new_config) + config = verify_and_sync_scan_layers(config) + + # Compute abstract model and logical-axis specs for downstream checkpoint alignment. + # We invoke create_nnx_abstract_model once (a lightweight abstract trace with no physical memory cost) + # both to initialize pure NNX sharded models and to cleanly extract the logical PartitionSpec tree + # (e.g., axis names like "kv_heads", "mlp_moe") required by _align_checkpoint_to_model_shapes. + _create_model, abstract_model = create_nnx_abstract_model( + config, mesh, devices, model_mode, rng_key, quant_mode_str=quant_mode_str + ) + _, _abs_state_for_specs = nnx.split(abstract_model) + specs = nnx.get_partition_spec(_abs_state_for_specs) + if config.pure_nnx: - _create_model, abstract_model = create_nnx_abstract_model( - config, mesh, devices, model_mode, rng_key, quant_mode_str=quant_mode_str - ) model = maxtext_utils_nnx.create_nnx_sharded_model(abstract_model, _create_model, mesh=mesh) # TODO: print debug_sharding info else: model = create_nnx_sharded_model_hybrid(config, mesh, devices, model_mode, rng_key) - # Compute logical-axis specs for downstream checkpoint alignment. - # The model-creation helpers above resolve specs internally for sharding, but - # the checkpoint-loading branch below needs the logical PartitionSpec tree - # (axis names like "kv_heads", "mlp_moe") for repeat/zero-pad dispatch in - # _align_checkpoint_to_model_shapes. nnx.eval_shape is cheap (abstract trace). - _create_model_for_specs = get_nnx_create_model_fn( - config, mesh, devices, model_mode, rng_key, quant_mode_str=quant_mode_str - ) - with nn.logical_axis_rules(config.logical_axis_rules): - _abs_model_for_specs = nnx.eval_shape(_create_model_for_specs) - _, _abs_state_for_specs = nnx.split(_abs_model_for_specs) - specs = nnx.get_partition_spec(_abs_state_for_specs) - sharded_state = nnx.state(model) if mesh is None: @@ -894,260 +934,269 @@ def from_pretrained( with mesh: if config.load_parameters_path: - with handle_checkpoint_mismatch("load parameters", config.load_parameters_path): - ckptr = ocp.Checkpointer( - ocp.PyTreeCheckpointHandler( - restore_concurrent_gb=config.checkpoint_storage_concurrent_gb, - save_concurrent_gb=config.checkpoint_storage_concurrent_gb, - use_ocdbt=config.checkpoint_storage_use_ocdbt, - use_zarr3=config.checkpoint_storage_use_zarr3, - ) - ) - - # This is a memory optimization. We don't want to restore the entire checkpoint - only the params. - # Rather than passing the entire abstract state, which could unnecessarily restore opt_state and - # waste memory, we instead restore the params field of the checkpoint (which itself may be a dictionary - # containing a key named 'params'). - - # Get the structure of checkpoint in `config.load_parameters_path` - metadata = ckptr.metadata(config.load_parameters_path) - if metadata is None or metadata.item_metadata is None: - max_logging.log( - f"ERROR: No valid Orbax checkpoint found at '{config.load_parameters_path}'. " - "Please check your load_parameters_path, the path may be missing, empty, " - "or point to a parent directory rather than the checkpoint step directory " - ) - raise ValueError( - f"No valid Orbax checkpoint found at '{config.load_parameters_path}'. " - "Please check your load_parameters_path." - ) - - def _adjust_target_for_moe_fusion(target, meta_tree, is_nnx): - if not hasattr(target, "items") or not hasattr(meta_tree, "items"): - return target - new_target = {} - for k, v in target.items(): - if k == "wi" and "wi" not in meta_tree and "wi_0" in meta_tree and "wi_1" in meta_tree: - if not is_nnx: - arr = v - half_dim = arr.shape[-1] // 2 - new_target["wi_0"] = jax.ShapeDtypeStruct( - shape=arr.shape[:-1] + (half_dim,), dtype=arr.dtype, sharding=arr.sharding - ) - new_target["wi_1"] = jax.ShapeDtypeStruct( - shape=arr.shape[:-1] + (half_dim,), dtype=arr.dtype, sharding=arr.sharding - ) - else: - arr = v["value"] - half_dim = arr.shape[-1] // 2 - new_target["wi_0"] = { - "value": jax.ShapeDtypeStruct( - shape=arr.shape[:-1] + (half_dim,), dtype=arr.dtype, sharding=arr.sharding - ) - } - new_target["wi_1"] = { - "value": jax.ShapeDtypeStruct( - shape=arr.shape[:-1] + (half_dim,), dtype=arr.dtype, sharding=arr.sharding - ) - } - else: - new_target[k] = _adjust_target_for_moe_fusion(v, meta_tree.get(k, {}), is_nnx) - - return new_target - - is_nnx_checkpoint = True - if ( - "params" in metadata.item_metadata.tree.keys() - and "params" in metadata.item_metadata.tree.get("params", {}).keys() - ): - # structure of linen checkpoint: {'params': {'params': {'decoder': ...}}} - is_nnx_checkpoint = False - target_for_restore = jax.tree.map( - lambda v: v[...], - sharded_state, - is_leaf=lambda n: isinstance(n, nnx.Variable), + ckptr = ocp.Checkpointer( + ocp.PyTreeCheckpointHandler( + restore_concurrent_gb=config.checkpoint_storage_concurrent_gb, + save_concurrent_gb=config.checkpoint_storage_concurrent_gb, + use_ocdbt=config.checkpoint_storage_use_ocdbt, + use_zarr3=config.checkpoint_storage_use_zarr3, ) + ) - target_for_restore = _adjust_target_for_moe_fusion( - target_for_restore, metadata.item_metadata.tree["params"]["params"], False - ) + # This is a memory optimization. We don't want to restore the entire checkpoint - only the params. + # Rather than passing the entire abstract state, which could unnecessarily restore opt_state and + # waste memory, we instead restore the params field of the checkpoint (which itself may be a dictionary + # containing a key named 'params'). + + # Get the structure of checkpoint in `config.load_parameters_path` + metadata = ckptr.metadata(config.load_parameters_path) + if metadata is None or metadata.item_metadata is None: + max_logging.log( + f"ERROR: No valid Orbax checkpoint found at '{config.load_parameters_path}'. " + "Please check your load_parameters_path, the path may be missing, empty, " + "or point to a parent directory rather than the checkpoint step directory " + ) + raise ValueError( + f"No valid Orbax checkpoint found at '{config.load_parameters_path}'. " + "Please check your load_parameters_path." + ) - item_to_restore = {"params": {"params": target_for_restore}} - base_restore_args = ocp.checkpoint_utils.construct_restore_args(target_for_restore) - restore_args = { - "params": { - "params": _fix_restore_args_for_shape_mismatch( - base_restore_args, - metadata.item_metadata.tree["params"]["params"], - mesh, + def _adjust_target_for_moe_fusion(target, meta_tree, is_nnx): + if not hasattr(target, "items") or not hasattr(meta_tree, "items"): + return target + new_target = {} + for k, v in target.items(): + if k == "wi" and "wi" not in meta_tree and "wi_0" in meta_tree and "wi_1" in meta_tree: + if not is_nnx: + arr = v + half_dim = arr.shape[-1] // 2 + new_target["wi_0"] = jax.ShapeDtypeStruct( + shape=arr.shape[:-1] + (half_dim,), dtype=arr.dtype, sharding=arr.sharding + ) + new_target["wi_1"] = jax.ShapeDtypeStruct( + shape=arr.shape[:-1] + (half_dim,), dtype=arr.dtype, sharding=arr.sharding + ) + else: + arr = v["value"] + half_dim = arr.shape[-1] // 2 + new_target["wi_0"] = { + "value": jax.ShapeDtypeStruct( + shape=arr.shape[:-1] + (half_dim,), dtype=arr.dtype, sharding=arr.sharding + ) + } + new_target["wi_1"] = { + "value": jax.ShapeDtypeStruct( + shape=arr.shape[:-1] + (half_dim,), dtype=arr.dtype, sharding=arr.sharding ) } - } - else: - # NNX checkpoint: {'decoder': {'value': ...}}, or NNX-RL with extra 'base' nesting. - # Restore only nnx.Param — RNG variable shapes may differ between checkpoint and model, - # and pure-dict checkpoints written by `layerwise_quantization._load_and_quantize_nnx` - # don't carry RNG/dropout state at all (they only persist nnx.Param leaves, including - # AQT serve-mode `qrhs.frozen` which is a Param subclass). - def _build_value_target(v): - # `v[...]` (a.k.a. `v.get_value(index=...)`) descends into the inner - # value with `value[Ellipsis]`. AQT serve-mode `qrhs.frozen` variables - # wrap a QTensor whose `__getitem__` calls `qvalue[idx]` on a - # `LogicallyPartitioned` wrapper — that fails. For QTensor (and any - # composite pytree value), use the unwrapped value directly so the - # restore target preserves the QTensor's qvalue/scale sub-structure. - inner = v.get_value() if hasattr(v, "get_value") else v[...] - if hasattr(inner, "shape"): - return {"value": v[...]} - # AQT QTensor: qvalue/scale leaves come back wrapped in flax - # `Partitioned` (a logical-axis sharding box). The on-disk save in - # `_load_and_quantize_nnx` flushes the QTensor as plain arrays — - # paths look like `qrhs.frozen.value.qvalue` / `...scale.0`. If we - # leave Partitioned in place, jax.tree adds an extra `.value` key - # under each leaf (`qrhs.frozen.value.qvalue.value`) and orbax - # silently fills with zeros because that path doesn't exist on - # disk. Strip Partitioned wrappers so the target tree matches. - inner = jax.tree.map( - lambda x: x.value if isinstance(x, Partitioned) else x, - inner, - is_leaf=lambda x: isinstance(x, Partitioned), - ) - return {"value": inner} - - # Keep persisted weight-like leaves: `nnx.Param` plus AQT serve-mode - # `qrhs.frozen` (a separate `aqt` Variable type, NOT a Param subclass). - # Excluded: `nnx.RngState` (regenerated per load, shapes can drift) and - # `nnx.Cache` (PREFILL/AR scratch, not persisted). Pure-dict checkpoints - # written by `layerwise_quantization._load_and_quantize_nnx` carry both - # Param kernels and `aqt`-typed `qrhs.frozen` quantized payloads. - if hasattr(sharded_state, "filter"): - param_state = sharded_state.filter(lambda path, var: not isinstance(var, (nnx.RngState, nnx.Cache))) else: - param_state = sharded_state - target_for_restore = jax.tree.map( - _build_value_target, - param_state, - is_leaf=lambda n: isinstance(n, nnx.Variable), - ) - has_base_key = "base" in metadata.item_metadata.tree - meta_tree_for_params = metadata.item_metadata.tree.get("base", metadata.item_metadata.tree) - target_for_restore = _adjust_target_for_moe_fusion(target_for_restore, meta_tree_for_params, True) - item_to_restore = {"base": target_for_restore} if has_base_key else target_for_restore - restore_args = _fix_restore_args_for_shape_mismatch( - ocp.checkpoint_utils.construct_restore_args(target_for_restore), meta_tree_for_params, mesh - ) - restore_args = {"base": restore_args} if has_base_key else restore_args - - # Free memory used by initial sharded_state before restore, to make room for the incoming checkpoint arrays. - # Skip nnx.Cache variables — they hold runtime state (e.g. GDN conv/recurrent state) that is - # not present in the checkpoint and must remain valid after the restore. - def _free_device_memory(node): - if isinstance(node, nnx.Variable) and not isinstance(node, (nnx.RngState, nnx.Cache)): - inner = node.get_value() if hasattr(node, "get_value") else node[...] - # AQT serve-mode `qrhs.frozen` wraps a QTensor (composite pytree) rather - # than a single jax.Array. Walking via tree_leaves frees the qvalue/scale - # arrays too; the single-leaf case is a 1-element tree. - for leaf in jax.tree_util.tree_leaves(inner): - if isinstance(leaf, jax.Array) and not leaf.is_deleted(): - leaf.delete() - elif isinstance(node, jax.Array) and not node.is_deleted(): - node.delete() - - return node - - jax.tree_util.tree_map(_free_device_memory, sharded_state, is_leaf=lambda n: isinstance(n, nnx.Variable)) - - restored = ckptr.restore( - epath.Path(config.load_parameters_path), - item=item_to_restore, - transforms={}, - restore_args=restore_args, + new_target[k] = _adjust_target_for_moe_fusion(v, meta_tree.get(k, {}), is_nnx) + + return new_target + + # Filter out transient runtime variables and rngs from NNX state before constructing target_for_restore. + # Checkpoints on disk only store persistent weight-like parameters. If we pass the full + # sharded_state directly, Orbax checks the disk for transient runtime state (like the layer + # dropout RNG seeds) and throws a structure mismatch error when it fails to find them. + # Note: We cannot use `nnx.split_state(sharded_state, nnx.Param)` to isolate weights because AQT + # serve-mode quantized models store their scale factors and integer payloads in custom AQT variable + # types (e.g. `qrhs.frozen`), which are NOT subclasses of `nnx.Param`. Negative filtering with + # `not isinstance(...)` safely retains all weight-like leaves while excluding transient runtime state. + param_state = sharded_state.filter( + lambda path, var: not isinstance(var, (nnx.RngState, nnx.Cache, nnx.Intermediate, nnx.BatchStat)) + ) + is_nnx_checkpoint = True + if ( + "params" in metadata.item_metadata.tree.keys() + and "params" in metadata.item_metadata.tree.get("params", {}).keys() + ): + # structure of linen checkpoint: {'params': {'params': {'decoder': ...}}} + is_nnx_checkpoint = False + target_for_restore = jax.tree.map( + lambda v: v[...], + param_state, + is_leaf=lambda n: isinstance(n, nnx.Variable), ) - if is_nnx_checkpoint: - restored_root = restored["base"] if has_base_key else restored - checkpoint = jax.tree.map( - lambda v: v["value"], - restored_root, - is_leaf=lambda x: isinstance(x, dict) and "value" in x and not isinstance(x.get("value"), dict), - ) - else: - checkpoint = restored["params"]["params"] - - if checkpoint: - # Same QTensor caveat as `_build_value_target` / `_free_device_memory`: - # `v[...]` fails on Variables wrapping QTensors. Use `get_value()` to - # access the inner value directly without index-style descent. - def _unwrap_for_align(v): - return v.get_value() if hasattr(v, "get_value") else v[...] - - model_arrays = jax.tree.map( - _unwrap_for_align, - sharded_state, - is_leaf=lambda n: isinstance(n, nnx.Variable), - ) - # ``specs`` (nnx.get_partition_spec(abstract_state) at the top of from_pretrained) - # is the source of truth for logical axis names — it's the input to - # nn.logical_to_mesh_sharding. Each leaf is a PartitionSpec whose entries are - # logical axis names (or None / nested tuples). Reuse it for repeat/zero-pad - # dispatch in _align_checkpoint_to_model_shapes. - # nnx.get_partition_spec returns Variables wrapping PartitionSpecs at the leaves; - # unwrap to raw PartitionSpecs so _normalize_logical_axes can read them. - logical_axes_tree = jax.tree.map( - lambda v: v.get_value(), - specs, - is_leaf=lambda n: isinstance(n, nnx.Variable), - ) + target_for_restore = _adjust_target_for_moe_fusion( + target_for_restore, metadata.item_metadata.tree["params"]["params"], False + ) - def to_dict(tree): - if hasattr(tree, "items"): - return {k: to_dict(v) for k, v in tree.items()} - return tree - - model_arrays = to_dict(model_arrays) - checkpoint = to_dict(checkpoint) - logical_axes_tree = to_dict(logical_axes_tree) - - checkpoint = _fuse_moe_weights(checkpoint, model_arrays) - # Release the raw restored buffers now that wi_0/wi_1 have been fused (if needed). - # This prevents the replicated intermediate copies from persisting until function return. - del restored - - def _filter_to_model_keys(ckpt, model): - """Recursively keep only keys present in model, dropping checkpoint-only fields (e.g. to_nnx__rngs).""" - if not hasattr(ckpt, "items") or not hasattr(model, "items"): - return ckpt - return {k: _filter_to_model_keys(ckpt[k], model[k]) for k in model if k in ckpt} - - checkpoint = _filter_to_model_keys(checkpoint, model_arrays) - - def _walk_align(ckpt, model_arr, axes): - if isinstance(ckpt, dict): - return { - k: _walk_align( - v, - model_arr[k], - axes.get(k) if isinstance(axes, dict) else None, - ) - for k, v in ckpt.items() - } - # AQT serve-mode `qrhs.frozen` wraps a QTensor (composite pytree of - # qvalue+scale arrays), not a single jax.Array. Shape alignment - # only makes sense for full-precision kernels — quantized payloads - # are saved in the exact shape the model expects, so pass through. - if not isinstance(ckpt, (jax.Array, jax.ShapeDtypeStruct, np.ndarray)): - return ckpt - return _align_checkpoint_to_model_shapes(ckpt, model_arr, axes) - - checkpoint = _walk_align(checkpoint, model_arrays, logical_axes_tree) - nnx.update(model, checkpoint) - else: - raise ValueError( - f"Checkpoint restore from '{config.load_parameters_path}' yielded no parameters. " - "This usually means the checkpoint format is incompatible with the model configuration " - "(e.g. a scanned checkpoint loaded with scan_layers=False, or vice versa). " - "Please ensure the checkpoint format matches the scan_layers setting." + item_to_restore = {"params": {"params": target_for_restore}} + base_restore_args = ocp.checkpoint_utils.construct_restore_args(target_for_restore) + restore_args = { + "params": { + "params": _fix_restore_args_for_shape_mismatch( + base_restore_args, + metadata.item_metadata.tree["params"]["params"], + mesh, + ) + } + } + else: + # NNX checkpoint: {'decoder': {'value': ...}}, or NNX-RL with extra 'base' nesting. + # Restore only nnx.Param — RNG variable shapes may differ between checkpoint and model, + # and pure-dict checkpoints written by `layerwise_quantization._load_and_quantize_nnx` + # don't carry RNG/dropout state at all (they only persist weight-like leaves, including + # AQT serve-mode `qrhs.frozen` which is not an aqt Variable subclass). + def _build_value_target(v): + # `v[...]` (a.k.a. `v.get_value(index=...)`) descends into the inner + # value with `value[Ellipsis]`. AQT serve-mode `qrhs.frozen` variables + # wrap a QTensor whose `__getitem__` calls `qvalue[idx]` on a + # `LogicallyPartitioned` wrapper — that fails. For QTensor (and any + # composite pytree value), use the unwrapped value directly so the + # restore target preserves the QTensor's qvalue/scale sub-structure. + inner = v.get_value() if hasattr(v, "get_value") else v[...] + if hasattr(inner, "shape"): + return {"value": v[...]} + # AQT QTensor: qvalue/scale leaves come back wrapped in flax + # `Partitioned` (a logical-axis sharding box). The on-disk save in + # `_load_and_quantize_nnx` flushes the QTensor as plain arrays — + # paths look like `qrhs.frozen.value.qvalue` / `...scale.0`. If we + # leave Partitioned in place, jax.tree adds an extra `.value` key + # under each leaf (`qrhs.frozen.value.qvalue.value`) and orbax + # silently fills with zeros because that path doesn't exist on + # disk. Strip Partitioned wrappers so the target tree matches. + inner = jax.tree.map( + lambda x: x.value if isinstance(x, Partitioned) else x, + inner, + is_leaf=lambda x: isinstance(x, Partitioned), ) + return {"value": inner} + + # Keep persisted weight-like leaves: `nnx.Param` plus AQT serve-mode + # `qrhs.frozen` (a separate `aqt` Variable type, NOT a Param subclass). + # Excluded: `nnx.RngState` (regenerated per load, shapes can drift), + # `nnx.Cache` (PREFILL/AR scratch, not persisted), `nnx.Intermediate` + # (transient forward-pass activations), and `nnx.BatchStat` (runtime statistics). + # Pure-dict checkpoints written by `layerwise_quantization._load_and_quantize_nnx` + # carry both Param kernels and `aqt`-typed `qrhs.frozen` quantized payloads. + target_for_restore = jax.tree.map( + _build_value_target, + param_state, + is_leaf=lambda n: isinstance(n, nnx.Variable), + ) + has_base_key = "base" in metadata.item_metadata.tree + meta_tree_for_params = metadata.item_metadata.tree.get("base", metadata.item_metadata.tree) + target_for_restore = _adjust_target_for_moe_fusion(target_for_restore, meta_tree_for_params, True) + item_to_restore = {"base": target_for_restore} if has_base_key else target_for_restore + restore_args = _fix_restore_args_for_shape_mismatch( + ocp.checkpoint_utils.construct_restore_args(target_for_restore), meta_tree_for_params, mesh + ) + restore_args = {"base": restore_args} if has_base_key else restore_args + + # Free memory used by initial sharded_state before restore, to make room for the incoming checkpoint arrays. + # Skip transient runtime variables (RngState, Cache, Intermediate, BatchStat) — they hold runtime state + # that is not present in the checkpoint and must remain valid after the restore. + def _free_device_memory(node): + if isinstance(node, nnx.Variable) and not isinstance( + node, (nnx.RngState, nnx.Cache, nnx.Intermediate, nnx.BatchStat) + ): + inner = node.get_value() if hasattr(node, "get_value") else node[...] + # AQT serve-mode `qrhs.frozen` wraps a QTensor (composite pytree) rather + # than a single jax.Array. Walking via tree_leaves frees the qvalue/scale + # arrays too; the single-leaf case is a 1-element tree. + for leaf in jax.tree_util.tree_leaves(inner): + if isinstance(leaf, jax.Array) and not leaf.is_deleted(): + leaf.delete() + elif isinstance(node, jax.Array) and not node.is_deleted(): + node.delete() + + return node + + jax.tree_util.tree_map(_free_device_memory, sharded_state, is_leaf=lambda n: isinstance(n, nnx.Variable)) + + restored = ckptr.restore( + epath.Path(config.load_parameters_path), + item=item_to_restore, + transforms={}, + restore_args=restore_args, + ) + + if is_nnx_checkpoint: + restored_root = restored["base"] if has_base_key else restored + checkpoint = jax.tree.map( + lambda v: v["value"], + restored_root, + is_leaf=lambda x: isinstance(x, dict) and "value" in x and not isinstance(x.get("value"), dict), + ) + else: + checkpoint = restored["params"]["params"] + + if checkpoint: + # Same QTensor caveat as `_build_value_target` / `_free_device_memory`: + # `v[...]` fails on Variables wrapping QTensors. Use `get_value()` to + # access the inner value directly without index-style descent. + def _unwrap_for_align(v): + return v.get_value() if hasattr(v, "get_value") else v[...] + + model_arrays = jax.tree.map( + _unwrap_for_align, + sharded_state, + is_leaf=lambda n: isinstance(n, nnx.Variable), + ) + # ``specs`` (nnx.get_partition_spec(abstract_state) at the top of from_pretrained) + # is the source of truth for logical axis names — it's the input to + # nn.logical_to_mesh_sharding. Each leaf is a PartitionSpec whose entries are + # logical axis names (or None / nested tuples). Reuse it for repeat/zero-pad + # dispatch in _align_checkpoint_to_model_shapes. + # nnx.get_partition_spec returns Variables wrapping PartitionSpecs at the leaves; + # unwrap to raw PartitionSpecs so _normalize_logical_axes can read them. + logical_axes_tree = jax.tree.map( + lambda v: v.get_value(), + specs, + is_leaf=lambda n: isinstance(n, nnx.Variable), + ) + + def to_dict(tree): + if hasattr(tree, "items"): + return {k: to_dict(v) for k, v in tree.items()} + return tree + + model_arrays = to_dict(model_arrays) + checkpoint = to_dict(checkpoint) + logical_axes_tree = to_dict(logical_axes_tree) + + checkpoint = _fuse_moe_weights(checkpoint, model_arrays) + # Release the raw restored buffers now that wi_0/wi_1 have been fused (if needed). + # This prevents the replicated intermediate copies from persisting until function return. + del restored + + def _filter_to_model_keys(ckpt, model): + """Recursively keep only keys present in model, dropping checkpoint-only fields (e.g. to_nnx__rngs).""" + if not hasattr(ckpt, "items") or not hasattr(model, "items"): + return ckpt + return {k: _filter_to_model_keys(ckpt[k], model[k]) for k in model if k in ckpt} + + checkpoint = _filter_to_model_keys(checkpoint, model_arrays) + + def _walk_align(ckpt, model_arr, axes): + if isinstance(ckpt, dict): + return { + k: _walk_align( + v, + model_arr[k], + axes.get(k) if isinstance(axes, dict) else None, + ) + for k, v in ckpt.items() + } + # AQT serve-mode `qrhs.frozen` wraps a QTensor (composite pytree of + # qvalue+scale arrays), not a single jax.Array. Shape alignment + # only makes sense for full-precision kernels — quantized payloads + # are saved in the exact shape the model expects, so pass through. + if not isinstance(ckpt, (jax.Array, jax.ShapeDtypeStruct, np.ndarray)): + return ckpt + return _align_checkpoint_to_model_shapes(ckpt, model_arr, axes) + + checkpoint = _walk_align(checkpoint, model_arrays, logical_axes_tree) + nnx.update(model, checkpoint) + else: + raise ValueError( + f"Checkpoint restore from '{config.load_parameters_path}' yielded no parameters. " + "This usually means the checkpoint format is incompatible with the model configuration " + "(e.g. a scanned checkpoint loaded with scan_layers=False, or vice versa). " + "Please ensure the checkpoint format matches the scan_layers setting." + ) if wrap_with_tunix_adapter: with mesh: diff --git a/src/maxtext/utils/qk_clip_utils.py b/src/maxtext/utils/qk_clip_utils.py index 87c3688fcf..7c26f130d4 100644 --- a/src/maxtext/utils/qk_clip_utils.py +++ b/src/maxtext/utils/qk_clip_utils.py @@ -153,7 +153,7 @@ def apply_qk_clip_nnx(state, intermediate_outputs, config): tau = float(config.qk_clip_threshold) _, params_state, _ = nnx.split(state.model, nnx.Param, ...) - params_dict = params_state.to_pure_dict() + params_dict = params_state.to_pure_dict() # pyrefly: ignore[missing-attribute] def clip_mla_weights(path, param): if len(path) < 2: diff --git a/src/maxtext/utils/sharding.py b/src/maxtext/utils/sharding.py index dc2f70ae14..3b524158cc 100644 --- a/src/maxtext/utils/sharding.py +++ b/src/maxtext/utils/sharding.py @@ -176,9 +176,7 @@ def remove_size_one_mesh_axis(spec, mesh): return P(*new_spec, unreduced=spec.unreduced, reduced=spec.reduced) -def get_nnx_var_named_sharding_with_scan_axis( - v: nnx.Variable, mesh -) -> nnx.Variable: +def get_nnx_var_named_sharding_with_scan_axis(v: nnx.Variable, mesh) -> nnx.Variable: """Compute NamedSharding for an NNX variable, correctly handling the scan axis.""" val = v.get_value() if not hasattr(val, "shape"): @@ -190,11 +188,7 @@ def get_nnx_var_named_sharding_with_scan_axis( return v.replace(jax.tree.map(lambda _: replicated, val)) return v metadata = v.get_metadata() - out_sharding = ( - metadata.get("out_sharding") - or metadata.get("sharding_names") - or metadata.get("sharding") - ) + out_sharding = metadata.get("out_sharding") or metadata.get("sharding_names") or metadata.get("sharding") if not out_sharding: pspec = P() else: @@ -202,11 +196,7 @@ def get_nnx_var_named_sharding_with_scan_axis( if nnx.PARTITION_NAME in metadata: partition_name = metadata[nnx.PARTITION_NAME] scan_axis = metadata.get("param_scan_axis", 0) - out_sharding = ( - [out_sharding] - if isinstance(out_sharding, str) - else list(out_sharding) - ) + out_sharding = [out_sharding] if isinstance(out_sharding, str) else list(out_sharding) if partition_name not in out_sharding: out_sharding.insert(scan_axis, partition_name) out_sharding = tuple(out_sharding) @@ -215,9 +205,7 @@ def get_nnx_var_named_sharding_with_scan_axis( local_rules = metadata.get("sharding_rules", ()) if context_rules or local_rules: local_rules_list = list(local_rules) if local_rules is not None else [] - context_rules_list = ( - list(context_rules) if context_rules is not None else [] - ) + context_rules_list = list(context_rules) if context_rules is not None else [] rules = local_rules_list + context_rules_list pspec = logical_to_mesh_axes(out_sharding, mesh, rules=rules) else: @@ -344,7 +332,6 @@ def _get_nontrival_mesh_axes(mesh): "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "stage", "expert", @@ -540,7 +527,7 @@ def add_data_to_sharding(mesh, path, aval, sharding): partition = (partition,) if size % mesh.shape["data"] == 0 and (partition is None or "tensor" not in partition): - added_component = ("data",) + partition + added_component = ("data",) + partition # pyrefly: ignore[unsupported-operation] new_pspec = jax.sharding.PartitionSpec(*(pspec[:idx] + (added_component,) + pspec[idx + 1 :])) new_sharding = jax.sharding.NamedSharding(sharding.mesh, new_pspec) return new_sharding @@ -581,7 +568,7 @@ def maybe_update_params_sharding_with_opt(config, state_mesh_shardings): # When quantization=fp8 is enabled the sharded_fp32_params # are not wrapped in `params`. Here we wrap them back. sharded_fp32_params = {"params": sharded_fp32_params} - state_mesh_shardings = state_mesh_shardings.replace(params=dict(prev_params_shardings, **sharded_fp32_params)) + state_mesh_shardings = state_mesh_shardings.replace(params=dict(prev_params_shardings, **sharded_fp32_params)) # pyrefly: ignore[bad-unpacking] return prev_params_shardings, state_mesh_shardings @@ -701,9 +688,7 @@ def _update_model_var(path, var): return prev_params_shardings, updated_state -def build_zero1_input_state_mesh_shardings( - config, state_mesh_shardings, params_shardings -): +def build_zero1_input_state_mesh_shardings(config, state_mesh_shardings, params_shardings): """Build the train-step input shardings under shard_optimizer_over_data (Zero-1). Model params on input use the original pre-Zero-1 sharding (params_shardings), diff --git a/src/maxtext/utils/standalone_checkpointer.py b/src/maxtext/utils/standalone_checkpointer.py index 6b1aa264c2..243ad1dd3e 100644 --- a/src/maxtext/utils/standalone_checkpointer.py +++ b/src/maxtext/utils/standalone_checkpointer.py @@ -25,7 +25,6 @@ from absl import app from flax import nnx -from flax.linen import partitioning as nn_partitioning import jax from jax import numpy as jnp from maxtext.configs import pyconfig @@ -69,36 +68,25 @@ def init_state_fn(): mesh = model.mesh _, tx = train_utils.create_training_optimizer(config, model) init_state_fn = partial(maxtext_utils.init_initial_state, model, tx, config, True, init_rng) + checkpoint_manager = train_utils.create_checkpoint_manager(config, mesh, init_state_fn) - unboxed_abstract_state, _, _ = maxtext_utils.get_abstract_state(config, mesh, init_state_fn, is_training=True) # A barrier to sync all hosts before starting to restore checkpoint jax.experimental.multihost_utils.sync_global_devices("Barrier before load") - checkpoint_load_start = datetime.datetime.now() - with nn_partitioning.axis_rules(config.logical_axis_rules): - state, _ = checkpointing.load_state_if_possible( - checkpoint_manager, - None, - config.load_parameters_path, - config.load_full_state_path, - config.checkpoint_storage_concurrent_gb, - unboxed_abstract_state, - use_ocdbt=config.checkpoint_storage_use_ocdbt, - use_zarr3=config.checkpoint_storage_use_zarr3, - ) - if state: - state = state["items"] + checkpoint_load_start = datetime.datetime.now() + # Delegate checkpoint restoration or state initialization to setup_training_state + state, _, _, _, was_restored = maxtext_utils.setup_training_state(None, config, mesh, checkpoint_manager, init_state_fn) jax.block_until_ready(state) checkpoint_load_end = datetime.datetime.now() - if state is not None: # Checkpoint was available for restore + + if was_restored: if jax.process_index() == 0: max_logging.log( "STANDALONE CHECKPOINTER : Checkpoint restored in :" f" {checkpoint_load_end - checkpoint_load_start}" ) - else: # Checkpoint was unavailable, state needs to be initialized - state, _, _, _ = maxtext_utils.setup_training_state(None, config, mesh, checkpoint_manager, init_state_fn) - state = add_entropy_to_checkpoint(state) + else: # Checkpoint was unavailable, fresh state needs to be perturbed with entropy + state = add_entropy_to_checkpoint(state) start_step = get_first_step(model, state) # this is the start_step for training for step in np.arange(start_step, config.steps): @@ -106,7 +94,8 @@ def init_state_fn(): start_time = datetime.datetime.now() # A barrier to sync all hosts before starting to save checkpoint jax.experimental.multihost_utils.sync_global_devices("Barrier before save") - if checkpointing.save_checkpoint(checkpoint_manager, int(step), state): + state_to_save = train_state_nnx.to_linen_checkpoint_dict(state.to_pure_dict()) if config.pure_nnx else state + if checkpointing.save_checkpoint(checkpoint_manager, int(step), state_to_save): checkpoint_manager.wait_until_finished() end_time = datetime.datetime.now() if jax.process_index() == 0: diff --git a/src/maxtext/utils/train_utils.py b/src/maxtext/utils/train_utils.py index 73e5d06b05..0b8fb6564e 100644 --- a/src/maxtext/utils/train_utils.py +++ b/src/maxtext/utils/train_utils.py @@ -299,7 +299,7 @@ def create_train_state_fn(): # Create data_loader AFTER reordering wrapper is applied data_loader = create_dataloader(config, mesh, data_iterator, recorder, rampup_manager) - state, _, state_mesh_shardings, data_iterator = maxtext_utils.setup_training_state( + state, _, state_mesh_shardings, data_iterator, _ = maxtext_utils.setup_training_state( data_iterator, config, mesh, checkpoint_manager, init_state_fn ) if config.pure_nnx: @@ -321,21 +321,17 @@ def create_train_state_fn(): state, outer_opt_state_sharding = diloco.build_diloco_state(config, lambda: state, mesh=mesh) # create state_mesh_shardings for the DilocoState - step_mesh = ( - state_mesh_shardings.optimizer.step.mesh - if config.pure_nnx - else state_mesh_shardings.step.mesh - ) + step_mesh = state_mesh_shardings.optimizer.step.mesh if config.pure_nnx else state_mesh_shardings.step.mesh inner_state_shardings = diloco.add_diloco_to_sharding(state_mesh_shardings) state_mesh_shardings = diloco.DiLoCoTrainState( inner_state_shardings, # Match the outer params' pure-dict structure (build_diloco_state stores # outer_params via to_pure_dict), so the sharding tree matches the state tree. - state_mesh_shardings_params.to_pure_dict() + state_mesh_shardings_params.to_pure_dict() # pyrefly: ignore[missing-attribute] if config.pure_nnx else state_mesh_shardings_params, outer_opt_state_sharding, - jax.sharding.NamedSharding( + jax.sharding.NamedSharding( # pyrefly: ignore[bad-argument-type] mesh=step_mesh, spec=jax.sharding.PartitionSpec() ), ) @@ -355,7 +351,7 @@ def create_train_state_fn(): logical_annotations = maxtext_utils.get_logical_annotations(config, mesh, init_state_fn) logical_annotations_params = logical_annotations.params - max_utils.print_non_trivial_mesh_axis(model.mesh) + max_utils.print_non_trivial_mesh_axis(model.mesh) # pyrefly: ignore[missing-attribute] maxtext_utils.print_shardings_params(state_params, state_mesh_shardings_params, mesh, logical_annotations_params) if config.pure_nnx: @@ -363,9 +359,9 @@ def create_train_state_fn(): # Don't merge the DiLoCoTrainState into the plain-model graphdef. The inner # train step needs that graphdef as jit_model; the wrapper passes through as state. train_state = state - model = state_graphdef + model = state_graphdef # pyrefly: ignore[unbound-name] else: - train_state = nnx.merge(state_graphdef, state) + train_state = nnx.merge(state_graphdef, state) # pyrefly: ignore[unbound-name] model = train_state.model else: train_state = state diff --git a/tests/assets/logits_generation/generate_hf_golden_logits.py b/tests/assets/logits_generation/generate_hf_golden_logits.py index c57d58c380..5b59214d62 100644 --- a/tests/assets/logits_generation/generate_hf_golden_logits.py +++ b/tests/assets/logits_generation/generate_hf_golden_logits.py @@ -80,6 +80,10 @@ def save_golden_logits( from transformers import Llama4ForConditionalGeneration # pylint: disable=import-outside-toplevel model_class = Llama4ForConditionalGeneration + elif "qwen3-vl" in model_id.lower(): + from transformers import Qwen3VLForConditionalGeneration # pylint: disable=import-outside-toplevel + + model_class = Qwen3VLForConditionalGeneration else: from transformers import AutoModelForCausalLM # pylint: disable=import-outside-toplevel @@ -151,7 +155,10 @@ def save_golden_logits( for key, value in inputs.items(): new_key = "tokens" if key == "input_ids" else key val_np = value.cpu().numpy() - data_to_save[new_key] = val_np[0] if val_np.ndim > 0 else val_np + if key == "pixel_values" and "qwen3-vl" in model_id.lower(): + data_to_save[new_key] = val_np + else: + data_to_save[new_key] = val_np[0] if val_np.ndim > 0 else val_np data_to_save["logits"] = logits[0] print(f"Token length is {len(data_to_save['tokens'])} for prompt: {prompt_text}") diff --git a/tests/conftest.py b/tests/conftest.py index eec4afa225..cfb66c419c 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -115,6 +115,7 @@ def _custom_iter(self): import maxtext # pylint: disable=unused-import from maxtext.common.gcloud_stub import is_decoupled +from tests.utils.newly_added_detection import get_changed_tests # Configure JAX to use unsafe_rbg PRNG implementation to match main scripts. if is_decoupled(): @@ -149,6 +150,15 @@ def pytest_collection_modifyitems(config, items): - Deselect tests marked as external_serving/training in decoupled mode. - Mark remaining tests with the `decoupled` marker when running decoupled. """ + + changed_tests = get_changed_tests() # set[(file_path, test_name)] + if changed_tests: + for item in items: + item_file = item.nodeid.split("::", 1)[0] + base_name = getattr(item, "originalname", item.name) + if (item_file, base_name) in changed_tests or (item_file, item.name) in changed_tests: + item.add_marker(pytest.mark.newly_added) + decoupled = is_decoupled() remaining = [] deselected = [] @@ -199,6 +209,7 @@ def pytest_configure(config): "external_training: goodput integrations", "decoupled: marked on tests that are not skipped due to GCP deps, when DECOUPLE_GCLOUD=TRUE", "skip_on_tpu7x: skip test if running on TPU7x platform", + "newly_added: newly introduced or modified tests in PRs, executed even if scheduled_only", ]: config.addinivalue_line("markers", m) diff --git a/tests/end_to_end/tpu/gemma3/4b/test_gemma3_lora.sh b/tests/end_to_end/tpu/gemma3/4b/test_gemma3_lora.sh new file mode 100755 index 0000000000..59408e8cdd --- /dev/null +++ b/tests/end_to_end/tpu/gemma3/4b/test_gemma3_lora.sh @@ -0,0 +1,83 @@ +#!/bin/bash + +# Validates the Gemma3-4B LoRA pipeline using a pre-converted MaxText checkpoint. + +# The flow of this script is as follows: +# 1. Run inference on the pre-converted checkpoint. +# 2. Run LoRA starting from the pre-converted checkpoint. +# 3. Run inference on the checkpoint produced by the LoRA run. +# 4. Convert the checkpoint produced by the LoRA run back to HuggingFace format. + +# Usage: +# export HF_TOKEN= +# export RUN_ID=$(date +%Y-%m-%d-%H-%M) +# bash test_gemma3_to_mt.sh $RUN_ID +# bash test_gemma3_lora.sh $RUN_ID + + +set -ex + +run_id=${1:-$(date +%Y-%m-%d-%H-%M)} +MODEL_NAME='gemma3-4b' + +# Non-Googlers please remember to point `BASE_OUTPUT_DIRECTORY` to the GCS paths where you have the scanned and unscanned checkpoints stored +BASE_OUTPUT_DIRECTORY=gs://runner-maxtext-logs/${MODEL_NAME} +UNSCANNED_CKPT_PATH=${BASE_OUTPUT_DIRECTORY}/to_maxtext/unscanned/${run_id}/0/items +SCANNED_CKPT_PATH=${BASE_OUTPUT_DIRECTORY}/to_maxtext/scanned/${run_id}/0/items + +# Step 1: Install torch +python3 -m pip install torch --index-url https://download.pytorch.org/whl/cpu + +# Step 2: Run inference on the original checkpoint converted from Hugging Face +python3 -m maxtext.inference.vllm_decode \ + model_name=${MODEL_NAME} \ + load_parameters_path=${UNSCANNED_CKPT_PATH} \ + vllm_hf_overrides='{architectures: ["MaxTextForCausalLM"]}' \ + hbm_utilization_vllm=0.6 \ + prompt="Suggest some famous landmarks in London." \ + use_chat_template=True \ + scan_layers=false + +# Step 3: Run LoRA on the converted checkpoint +python3 -m maxtext.trainers.post_train.sft.train_sft \ + base_output_directory=${BASE_OUTPUT_DIRECTORY}/lora \ + load_parameters_path=${SCANNED_CKPT_PATH} \ + per_device_batch_size=1 run_name=${run_id} \ + steps=5 scan_layers=true \ + model_name=${MODEL_NAME} \ + learning_rate=3e-6 \ + lora.enable_lora=True \ + lora.lora_rank=16 \ + lora.lora_alpha=32.0 \ + enable_nnx=True \ + pure_nnx_decoder=True \ + enable_single_controller=True \ + checkpoint_storage_use_zarr3=False checkpoint_storage_use_ocdbt=False + + +# Step 4: Run inference on the checkpoint generated from the previous run +python3 -m maxtext.inference.vllm_decode \ + --use_tunix=True \ + model_name=${MODEL_NAME} \ + load_parameters_path=${SCANNED_CKPT_PATH} \ + lora.enable_lora=True \ + lora.lora_restore_path=${BASE_OUTPUT_DIRECTORY}/lora/${run_id}/checkpoints/5/model_params \ + lora.lora_rank=16 \ + lora.lora_alpha=32.0 \ + vllm_hf_overrides='{architectures: ["MaxTextForCausalLM"]}' \ + hbm_utilization_vllm=0.6 \ + prompt="Suggest some famous landmarks in London." \ + use_chat_template=True \ + enable_nnx=True \ + pure_nnx_decoder=True \ + scan_layers=true + +# Step 5: Convert the checkpoint from MaxText format to Hugging Face format +python3 -m maxtext.checkpoint_conversion.to_huggingface \ + model_name=${MODEL_NAME} \ + load_parameters_path=${SCANNED_CKPT_PATH} \ + lora.lora_restore_path=${BASE_OUTPUT_DIRECTORY}/lora/${run_id}/checkpoints/5/model_params \ + base_output_directory=${BASE_OUTPUT_DIRECTORY}/to_huggingface/unscanned/${run_id} \ + scan_layers=true \ + enable_nnx=True \ + pure_nnx_decoder=True diff --git a/tests/end_to_end/tpu/qwen3/30b/test_qwen3.sh b/tests/end_to_end/tpu/qwen3/30b/test_qwen3.sh new file mode 100755 index 0000000000..1f5c8f7a87 --- /dev/null +++ b/tests/end_to_end/tpu/qwen3/30b/test_qwen3.sh @@ -0,0 +1,58 @@ +#!/bin/bash + +# Validates the Qwen3-30B pre-training pipeline using a pre-converted MaxText checkpoint. + +# The flow of this script is as follows: +# 1. Run inference on the pre-converted checkpoint. +# 2. Run pre-training starting from the pre-converted checkpoint. +# 3. Run inference on the checkpoint produced by the pre-training run. + +# Usage: +# export HF_TOKEN= +# export RUN_ID=$(date +%Y-%m-%d-%H-%M-%S) +# bash test_qwen3_to_mt.sh $RUN_ID +# bash test_qwen3.sh $RUN_ID + + +set -ex + +run_id=${1:-$(date +%Y-%m-%d-%H-%M-%S)} +MODEL_NAME='qwen3-30b-a3b-base' + + +BASE_OUTPUT_DIRECTORY=gs://runner-maxtext-logs/${MODEL_NAME} +UNSCANNED_CKPT_PATH=${BASE_OUTPUT_DIRECTORY}/to_maxtext/unscanned/${run_id}/0/items +SCANNED_CKPT_PATH=${BASE_OUTPUT_DIRECTORY}/to_maxtext/scanned/${run_id}/0/items +export PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python + +DATASET_PATH=gs://maxtext-dataset + +# Step 1: Run inference on the pre-converted checkpoint +python3 -m maxtext.inference.decode \ + model_name=${MODEL_NAME} tokenizer_type="huggingface" \ + load_parameters_path=${UNSCANNED_CKPT_PATH} \ + per_device_batch_size=1 run_name=${run_id} \ + max_prefill_predict_length=8 max_target_length=16 steps=1 async_checkpointing=false \ + checkpoint_storage_use_zarr3=False checkpoint_storage_use_ocdbt=False \ + scan_layers=false prompt='I love to' attention=\'dot_product\' + +# Step 2: Run pre-training starting from the pre-converted checkpoint +python3 -m maxtext.trainers.pre_train.train \ + base_output_directory=${BASE_OUTPUT_DIRECTORY}/train \ + dataset_path=${DATASET_PATH} tokenizer_type="huggingface" \ + load_parameters_path=${SCANNED_CKPT_PATH} \ + per_device_batch_size=0.125 run_name=${run_id} \ + max_target_length=64 steps=5 async_checkpointing=false \ + checkpoint_storage_use_zarr3=False checkpoint_storage_use_ocdbt=False \ + model_name=${MODEL_NAME} scan_layers=true \ + remat_policy=full \ + ici_tensor_parallelism=4 ici_fsdp_parallelism=2 weight_dtype=bfloat16 dtype=bfloat16 opt_type=sgd optimizer_memory_host_offload=true + +# Step 3: Run inference on the checkpoint produced by the pre-training run +python3 -m maxtext.inference.decode \ + model_name=${MODEL_NAME} tokenizer_type="huggingface" \ + load_parameters_path=${BASE_OUTPUT_DIRECTORY}/train/${run_id}/checkpoints/4/items \ + per_device_batch_size=1 run_name=${run_id} \ + max_prefill_predict_length=8 max_target_length=16 steps=4 async_checkpointing=false \ + checkpoint_storage_use_zarr3=False checkpoint_storage_use_ocdbt=False \ + scan_layers=true prompt='I love to' attention=\'dot_product\' diff --git a/tests/end_to_end/tpu/qwen3/30b/test_qwen3_rl.sh b/tests/end_to_end/tpu/qwen3/30b/test_qwen3_rl.sh new file mode 100755 index 0000000000..6ac095cdf5 --- /dev/null +++ b/tests/end_to_end/tpu/qwen3/30b/test_qwen3_rl.sh @@ -0,0 +1,84 @@ +#!/bin/bash + +# Validates the Qwen3 RL pipeline using a pre-converted MaxText checkpoint. + +# The flow of this script is as follows: +# 1. Run inference on the pre-converted checkpoint. +# 2. Run RL starting from the pre-converted checkpoint. +# 3. Run inference on the checkpoint produced by the RL run. + +# Usage: +# export HF_TOKEN= +# export RUN_ID=$(date +%Y-%m-%d-%H-%M-%S) +# bash test_qwen3_to_mt.sh $RUN_ID +# bash test_qwen3_rl.sh $RUN_ID + +set -ex + +export VLLM_WORKER_MULTIPROC_METHOD=spawn +export PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=upb +export VLLM_RAY_EXTRA_ENV_VARS_TO_COPY="PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION" +export VLLM_ENABLE_V1_MULTIPROCESSING=0 +export JAX_RANDOM_WEIGHTS=1 +export SKIP_JAX_PRECOMPILE=1 +export NEW_MODEL_DESIGN=1 +export TPU_MIN_LOG_LEVEL=0 +export TF_CPP_MIN_LOG_LEVEL=0 +export TPU_STDERR_LOG_LEVEL=0 + +# Force Python to use "spawn" instead of "fork" for multiprocessing to prevent gRPC socket corruption +cat << 'EOF' > usercustomize.py +import multiprocessing +try: + multiprocessing.set_start_method("spawn", force=True) +except Exception: + pass +EOF +export PYTHONPATH=${PYTHONPATH:-.}:$(pwd) +run_id=${1:-$(date +%Y-%m-%d-%H-%M-%S)} +MODEL_NAME='qwen3-30b-a3b-base' + +BASE_OUTPUT_DIRECTORY=gs://runner-maxtext-logs/${MODEL_NAME} +UNSCANNED_CKPT_PATH=${BASE_OUTPUT_DIRECTORY}/to_maxtext/unscanned/${run_id}/0/items +SCANNED_CKPT_PATH=${BASE_OUTPUT_DIRECTORY}/to_maxtext/scanned/${run_id}/0/items + +# Step 1: Run inference on the pre-converted checkpoint +python3 -m maxtext.inference.vllm_decode \ + model_name=${MODEL_NAME} \ + load_parameters_path=${SCANNED_CKPT_PATH} \ + vllm_hf_overrides='{architectures: ["MaxTextForCausalLM"]}' \ + hbm_utilization_vllm=0.85 \ + prompt='Suggest some famous landmarks in London.' \ + max_target_length=256 max_num_batched_tokens=256 \ + ici_tensor_parallelism=8 \ + allow_split_physical_axes=True prefuse_moe_weights=True \ + use_chat_template=True scan_layers=True enable_single_controller=True + +# Step 2: Run RL starting from the pre-converted checkpoint +python3 -m maxtext.trainers.post_train.rl.train_rl \ + base_output_directory=${BASE_OUTPUT_DIRECTORY}/rl \ + load_parameters_path=${SCANNED_CKPT_PATH} \ + run_name=${run_id} rl.loss_algo='grpo' scan_layers=True \ + num_batches=5 batch_size=4 train_micro_batch_size=1 num_test_batches=5 \ + model_name=${MODEL_NAME} enable_single_controller=True \ + checkpoint_storage_use_zarr3=False checkpoint_storage_use_ocdbt=False \ + rollout_tensor_parallelism=8 \ + vllm_hf_overrides='{architectures: ["MaxTextForCausalLM"]}' \ + vllm_additional_config='{"maxtext_config": {"model_name": "'${MODEL_NAME}'", "allow_split_physical_axes": true, "scan_layers": false, "prefuse_moe_weights": True}}' \ + remat_policy=full hbm_utilization_vllm=0.55 use_pathways=True \ + chips_per_vm=8 ici_tensor_parallelism=4\ + max_target_length=512 weight_dtype=bfloat16 dtype=bfloat16 opt_type=sgd \ + enable_tunix_perf_metrics=True rl.num_generations=16 \ + debug.rl=False rl.reshard_chunk_size=1 + +# Step 3: Run inference on the checkpoint produced by the RL run +python3 -m maxtext.inference.vllm_decode \ + model_name=${MODEL_NAME} \ + load_parameters_path=${BASE_OUTPUT_DIRECTORY}/rl/${run_id}/checkpoints/actor/4/items \ + vllm_hf_overrides='{architectures: ["MaxTextForCausalLM"]}' \ + hbm_utilization_vllm=0.85 \ + prompt='Suggest some famous landmarks in London.' \ + max_target_length=256 max_num_batched_tokens=256 \ + ici_tensor_parallelism=8 \ + allow_split_physical_axes=True prefuse_moe_weights=True \ + use_chat_template=True scan_layers=True enable_single_controller=True diff --git a/tests/end_to_end/tpu/qwen3/30b/test_qwen3_sft.sh b/tests/end_to_end/tpu/qwen3/30b/test_qwen3_sft.sh new file mode 100755 index 0000000000..1b1e77500a --- /dev/null +++ b/tests/end_to_end/tpu/qwen3/30b/test_qwen3_sft.sh @@ -0,0 +1,56 @@ +#!/bin/bash + +# Validates the Qwen3 SFT pipeline using a pre-converted MaxText checkpoint. + + +# The flow of this script is as follows: +# 1. Run inference on the pre-converted checkpoint. +# 2. Run SFT starting from the pre-converted checkpoint. +# 3. Run inference on the checkpoint produced by the SFT run. + +# Usage: +# export HF_TOKEN= +# export RUN_ID=$(date +%Y-%m-%d-%H-%M-%S) +# bash test_qwen3_to_mt.sh $RUN_ID +# bash test_qwen3_sft.sh $RUN_ID + + +set -ex +export PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python +export NEW_MODEL_DESIGN=1 +run_id=${1:-$(date +%Y-%m-%d-%H-%M-%S)} +MODEL_NAME='qwen3-30b-a3b-base' + +BASE_OUTPUT_DIRECTORY=gs://runner-maxtext-logs/${MODEL_NAME} +UNSCANNED_CKPT_PATH=${BASE_OUTPUT_DIRECTORY}/to_maxtext/unscanned/${run_id}/0/items +SCANNED_CKPT_PATH=${BASE_OUTPUT_DIRECTORY}/to_maxtext/scanned/${run_id}/0/items + +# Step 1: Run inference on the pre-converted checkpoint +python3 -m maxtext.inference.vllm_decode \ + model_name=${MODEL_NAME} \ + load_parameters_path=${SCANNED_CKPT_PATH} \ + vllm_hf_overrides='{architectures: ["MaxTextForCausalLM"]}' \ + hbm_utilization_vllm=0.85 \ + use_chat_template=True scan_layers=true enable_single_controller=True \ + ici_tensor_parallelism=4 prompt="Suggest some famous landmarks in London." + +# Step 2: Run SFT starting from the pre-converted checkpoint +python3 -m maxtext.trainers.post_train.sft.train_sft \ + base_output_directory=${BASE_OUTPUT_DIRECTORY}/sft \ + load_parameters_path=${SCANNED_CKPT_PATH} \ + per_device_batch_size=0.125 run_name=${run_id} \ + steps=5 scan_layers=true \ + model_name=${MODEL_NAME} \ + checkpoint_storage_use_zarr3=False checkpoint_storage_use_ocdbt=False \ + remat_policy=full \ + ici_tensor_parallelism=1 ici_fsdp_parallelism=1 ici_expert_parallelism=8 \ + enable_single_controller=True max_target_length=16 weight_dtype=bfloat16 dtype=bfloat16 opt_type=sgd + +# Step 3: Run inference on the checkpoint produced by the SFT run +python3 -m maxtext.inference.vllm_decode \ + model_name=${MODEL_NAME} \ + load_parameters_path=${BASE_OUTPUT_DIRECTORY}/sft/${run_id}/checkpoints/5/model_params \ + vllm_hf_overrides='{architectures: ["MaxTextForCausalLM"]}' \ + hbm_utilization_vllm=0.85 \ + use_chat_template=True scan_layers=true enable_single_controller=True \ + ici_tensor_parallelism=4 prompt="Suggest some famous landmarks in London." diff --git a/tests/end_to_end/tpu/qwen3/30b/test_qwen3_to_hf.sh b/tests/end_to_end/tpu/qwen3/30b/test_qwen3_to_hf.sh new file mode 100755 index 0000000000..9c35e55a7d --- /dev/null +++ b/tests/end_to_end/tpu/qwen3/30b/test_qwen3_to_hf.sh @@ -0,0 +1,31 @@ +#!/bin/bash + +# Converts a MaxText checkpoint to a Hugging Face model checkpoint. + +# Usage: +# export RUN_ID=$(date +%Y-%m-%d-%H-%M-%S) +# bash test_qwen3_to_hf.sh $RUN_ID $CHECKPOINT_PATH $SCAN_LAYERS + + +set -ex + +run_id=$1 +CKPT_PATH=$2 +SCAN_LAYERS=${3:-false} + +MODEL_NAME='qwen3-30b-a3b-base' +BASE_OUTPUT_DIRECTORY="gs://runner-maxtext-logs/${MODEL_NAME}" + +if [ "${SCAN_LAYERS,,}" = "true" ]; then + scan_status="scanned" +else + scan_status="unscanned" +fi + +# Convert the checkpoint to Hugging Face format +python3 -m maxtext.checkpoint_conversion.to_huggingface \ + model_name=${MODEL_NAME} \ + tokenizer_type="huggingface" \ + load_parameters_path=${CKPT_PATH} \ + base_output_directory=${BASE_OUTPUT_DIRECTORY}/to_huggingface/${scan_status}/${run_id} \ + scan_layers=$SCAN_LAYERS diff --git a/tests/end_to_end/tpu/qwen3/30b/test_qwen3_to_mt.sh b/tests/end_to_end/tpu/qwen3/30b/test_qwen3_to_mt.sh new file mode 100755 index 0000000000..aa723f194c --- /dev/null +++ b/tests/end_to_end/tpu/qwen3/30b/test_qwen3_to_mt.sh @@ -0,0 +1,61 @@ +#!/bin/bash + +# Converts Qwen3-30B HuggingFace checkpoint to MaxText format and validates logit correctness. + +# The flow of this script is as follows: +# 1. Install PyTorch (CPU) required for checkpoint conversion. +# 2. Convert the HuggingFace checkpoint to MaxText format in both unscanned and scanned formats. +# 3. Run a forward pass logits check to verify the converted checkpoint matches the original HF model. + +# Usage: +# export HF_TOKEN= +# export RUN_ID=$(date +%Y-%m-%d-%H-%M-%S) +# bash test_qwen3_to_mt.sh $RUN_ID + +set -ex + +run_id=${1:-$(date +%Y-%m-%d-%H-%M-%S)} +MODEL_NAME='qwen3-30b-a3b-base' +HF_GOLDEN_MODEL='Qwen/Qwen3-30B-A3B-Base' + + +BASE_OUTPUT_DIRECTORY=gs://runner-maxtext-logs/${MODEL_NAME}/to_maxtext + +# Step 1: Install torch +python3 -m pip install torch --index-url https://download.pytorch.org/whl/cpu + +# Step 2: Convert the checkpoint from Hugging Face +python3 -m maxtext.checkpoint_conversion.to_maxtext \ + model_name=${MODEL_NAME} \ + base_output_directory=${BASE_OUTPUT_DIRECTORY}/unscanned/${run_id} \ + scan_layers=false \ + hardware=cpu skip_jax_distributed_system=True \ + checkpoint_storage_use_zarr3=False checkpoint_storage_use_ocdbt=False \ + --lazy_load_tensors=False \ + --eager_load_method='safetensors' + +UNSCANNED_CKPT_PATH=${BASE_OUTPUT_DIRECTORY}/unscanned/${run_id}/0/items +echo "Unscanned checkpoint path: ${UNSCANNED_CKPT_PATH}" + +# Convert to scanned format +python3 -m maxtext.checkpoint_conversion.to_maxtext \ + model_name=${MODEL_NAME} \ + base_output_directory=${BASE_OUTPUT_DIRECTORY}/scanned/${run_id} \ + scan_layers=true \ + hardware=cpu skip_jax_distributed_system=True \ + checkpoint_storage_use_zarr3=False checkpoint_storage_use_ocdbt=False \ + --lazy_load_tensors=False \ + --eager_load_method='safetensors' + +SCANNED_CKPT_PATH=${BASE_OUTPUT_DIRECTORY}/scanned/${run_id}/0/items +echo "Scanned checkpoint path: ${SCANNED_CKPT_PATH}" + +# Step 3: Run forward pass logits check +python3 -m tests.utils.forward_pass_logit_checker \ + load_parameters_path=${UNSCANNED_CKPT_PATH} \ + model_name=${MODEL_NAME} \ + scan_layers=false \ + --hf_model_path=${HF_GOLDEN_MODEL} \ + --max_kl_div=0.03 \ + --run_hf_model=true \ + hardware=cpu skip_jax_distributed_system=True diff --git a/tests/end_to_end/tpu/qwen3/vl_2b/test_qwen3_vl_2b_to_hf_e2e.sh b/tests/end_to_end/tpu/qwen3/vl_2b/test_qwen3_vl_2b_to_hf_e2e.sh new file mode 100644 index 0000000000..56a8e12073 --- /dev/null +++ b/tests/end_to_end/tpu/qwen3/vl_2b/test_qwen3_vl_2b_to_hf_e2e.sh @@ -0,0 +1,124 @@ +#!/bin/bash + +# This script is both an end-to-end test and documentation for converting a +# Qwen3-VL-2B MaxText checkpoint to Hugging Face format. Can be run on a v4-8. + +# The flow of this script is as follows: +# 1. Convert an original Hugging Face model checkpoint to MaxText format. +# 2. Run a forward pass check to compare the logits and KL divergence between +# the MaxText checkpoint and the Hugging Face checkpoint. +# 3. (Optional) Convert the resulting MaxText checkpoint back to Hugging Face format. + +# Pre-requisites: +# 1. Set HF_TOKEN environment variable to your Hugging Face access token. +# export HF_TOKEN= +# 2. Configure USE_MULTIMODAL (true for multimodal, false for text-only). +# 3. Configure USE_SCAN_LAYERS (true if checkpoint was trained with scanned layers, false otherwise). + +set -ex + + +MODEL_NAME='qwen3-vl-2b' +export MODEL_VARIATION='vl_2b' +export HF_MODEL=Qwen/Qwen3-VL-2B-Instruct + +idx=$(date +%Y-%m-%d-%H-%M) + +# Set USE_SCAN_LAYERS=true if the checkpoint was trained with scanned layers +USE_SCAN_LAYERS=false +if ${USE_SCAN_LAYERS}; then export CHECKPOINT_TYPE=scanned; else export CHECKPOINT_TYPE=unscanned; fi + +USE_MULTIMODAL=true + + +export HF_TOKEN= + +export MODEL_BUCKET= + +# Path to Maxtext converted to HF checkpoint +export LOCAL_PATH=/hf/${MODEL_NAME}/${idx} + + +# Installing torch for deps in forward_pass_logit_checker.py +python3 -m pip install torch --index-url https://download.pytorch.org/whl/cpu +python3 -m pip install decord + +# Check point conversion +python3 -m maxtext.checkpoint_conversion.to_maxtext \ + "${MAXTEXT_CONFIGS_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/maxtext/configs}"/base.yml \ + model_name=${MODEL_NAME} \ + base_output_directory=${MODEL_BUCKET}/${MODEL_NAME}/${CHECKPOINT_TYPE}/${idx} \ + scan_layers=false \ + hf_access_token=${HF_TOKEN} \ + weight_dtype=bfloat16 \ + hardware=cpu \ + skip_jax_distributed_system=True \ + checkpoint_storage_use_ocdbt=False \ + checkpoint_storage_use_zarr3=False \ + --eager_load_method=safetensors \ + --lazy_load_tensors=False + +# Path to MaxText checkpoint +export CKPT_PATH=${MODEL_BUCKET}/${MODEL_NAME}/${CHECKPOINT_TYPE}/${idx}/0/items + + +# Run forward pass logit checker to validate the converted checkpoint. +if [ "${USE_MULTIMODAL}" == true ]; then + TEST_PROMPT='Describe this image' + TEST_IMAGE='tests/assets/test_image.jpg' + export GOLDEN_LOGITS_PATH=/tmp/golden_qwen3_vl_2b_vision.jsonl + + python3 -m tests.assets.logits_generation.generate_hf_golden_logits \ + --model-id=${HF_MODEL} \ + --output-path=${GOLDEN_LOGITS_PATH} \ + --prompts="${TEST_PROMPT}" \ + --image-paths=${TEST_IMAGE} \ + --hf-model-path=${HF_MODEL} \ + --apply-chat-template \ + --output-format=json + + echo "=== Running MaxText Forward Pass Logit Checker ===" + python3 -m tests.utils.forward_pass_logit_checker \ + "${MAXTEXT_CONFIGS_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/maxtext/configs}"/base.yml \ + tokenizer_path=${HF_MODEL} \ + load_parameters_path=${CKPT_PATH} \ + model_name=${MODEL_NAME} \ + use_multimodal=${USE_MULTIMODAL} \ + scan_layers=${USE_SCAN_LAYERS} \ + dtype=float32 \ + matmul_precision=highest \ + per_device_batch_size=1 \ + attention=dot_product \ + prompt="${TEST_PROMPT}" \ + image_path=${TEST_IMAGE} \ + --max_kl_div=0.1 \ + --golden_logits_path=${GOLDEN_LOGITS_PATH} \ + override_model_config=true +else + echo "=== Running MaxText Forward Pass Logit Checker ===" + python3 -m tests.utils.forward_pass_logit_checker \ + "${MAXTEXT_CONFIGS_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/maxtext/configs}"/base.yml \ + tokenizer_path=${HF_MODEL} \ + load_parameters_path=${CKPT_PATH} \ + model_name=${MODEL_NAME} \ + use_multimodal=${USE_MULTIMODAL} \ + scan_layers=${USE_SCAN_LAYERS} \ + per_device_batch_size=1 \ + dtype=float32 \ + --max_kl_div=0.1 \ + --run_hf_model=true \ + --hf_model_path=${HF_MODEL} \ + override_model_config=true +fi + +# Optional: Convert checkpoint back to HF format +python3 -m maxtext.checkpoint_conversion.to_huggingface \ + "${MAXTEXT_CONFIGS_DIR:-${MAXTEXT_REPO_ROOT:-$PWD}/src/maxtext/configs}"/base.yml \ + model_name=${MODEL_NAME} \ + hf_access_token=${HF_TOKEN} \ + load_parameters_path=${CKPT_PATH} \ + base_output_directory=${LOCAL_PATH} \ + use_multimodal=${USE_MULTIMODAL} \ + scan_layers=${USE_SCAN_LAYERS} \ + override_model_config=true + diff --git a/tests/inference/kvcache_test.py b/tests/inference/kvcache_test.py index ba068e0607..93ff926713 100644 --- a/tests/inference/kvcache_test.py +++ b/tests/inference/kvcache_test.py @@ -29,7 +29,7 @@ class MlaKVCacheTest(unittest.TestCase): def setUp(self): super().setUp() self.rng = jax.random.PRNGKey(42) - self.batchsize = 8 + self.batchsize = 4 self.prefill_len = 100 self.target_len = 196 self.dtype = jnp.bfloat16 diff --git a/tests/integration/checkpointing_test.py b/tests/integration/checkpointing_test.py index fcd59687e9..461123f3a8 100644 --- a/tests/integration/checkpointing_test.py +++ b/tests/integration/checkpointing_test.py @@ -223,12 +223,10 @@ def get_cmd(steps, metrics_file): train_main(get_cmd(steps=1, metrics_file="saved_metrics_mismatch.txt")) # 2. Attempt to restore with scan_layers=False and assert ValueError - mismatch_command = get_cmd( - steps=2, metrics_file="restored_metrics_mismatch.txt" - ) + ["scan_layers=False"] + mismatch_command = get_cmd(steps=2, metrics_file="restored_metrics_mismatch.txt") + ["scan_layers=False"] with pytest.raises(ValueError) as excinfo: train_main(mismatch_command) - assert "Failed to restore checkpoint" in str(excinfo.value) + assert "Configuration mismatch" in str(excinfo.value) or "Failed to restore checkpoint" in str(excinfo.value) assert "scan_layers" in str(excinfo.value) diff --git a/tests/integration/deepseek_scan_engram_test.py b/tests/integration/deepseek_scan_engram_test.py index 63d97ef80c..6891fd83e2 100644 --- a/tests/integration/deepseek_scan_engram_test.py +++ b/tests/integration/deepseek_scan_engram_test.py @@ -127,9 +127,6 @@ def batch_decode(self, token_ids, *args, **kwargs): ] ) - if base_num_decoder_layers >= 4 and jax.device_count() <= 8 and any("TPU7x" in d.device_kind for d in jax.devices()): - pytest.skip("Compiling 4+ DeepSeek MoE/Engram scanned layers requires >= 8" " physical TPU chips on TPU7x") - devices_array = maxtext_utils.create_device_mesh(config) mesh = Mesh(devices_array, config.mesh_axes) @@ -176,16 +173,14 @@ def test_decoder_init_engram_2_8(self, mock_from_pretrained): """Test engram layers at indices 2 and 8.""" self._test_engram_pattern( mock_from_pretrained, - "1,4", + "1", [ "dense_layers_0_0", "dense_layers_engram_1", "dense_layers_2_2", - "moe_layers_3_3", - "moe_layers_engram_4", ], first_num_dense_layers=3, - base_num_decoder_layers=5, + base_num_decoder_layers=3, ) @pytest.mark.tpu_only @@ -206,8 +201,8 @@ def test_decoder_init_engram_4_9(self, mock_from_pretrained): """Test engram layers at indices 4 and 9 - last engram layer of block.""" self._test_engram_pattern( mock_from_pretrained, - "1,3", - ["dense_layers_0_0", "dense_layers_engram_1", "moe_layers_2_2", "moe_layers_engram_3"], + "1,2", + ["dense_layers_0_0", "dense_layers_engram_1", "moe_layers_engram_2"], first_num_dense_layers=2, - base_num_decoder_layers=4, + base_num_decoder_layers=3, ) diff --git a/tests/integration/hlo_diff_test.py b/tests/integration/hlo_diff_test.py index cc18713749..baeb8ebbe7 100644 --- a/tests/integration/hlo_diff_test.py +++ b/tests/integration/hlo_diff_test.py @@ -93,11 +93,11 @@ def get_filtered_lines(file_path): return False @pytest.mark.parametrize( - "test_id, config_file, overrides", + "test_id, model_name, overrides", [ ( "deepseek3", - "src/maxtext/configs/models/deepseek3-test.yml", + "deepseek3-test", { "compile_topology": "v6e-4", "base_num_decoder_layers": 4, @@ -107,7 +107,7 @@ def get_filtered_lines(file_path): ), ( "llama3_8b", - "src/maxtext/configs/models/llama3-8b.yml", + "llama3-8b", { "compile_topology": "v6e-4", "base_num_decoder_layers": 4, @@ -117,7 +117,7 @@ def get_filtered_lines(file_path): ), ( "qwen3_1.7b", - "src/maxtext/configs/models/qwen3-1.7b.yml", + "qwen3-1.7b", { "compile_topology": "v6e-4", "base_num_decoder_layers": 4, @@ -127,7 +127,7 @@ def get_filtered_lines(file_path): ), ], ) - def test_hlo_diff(self, test_id, config_file, overrides): + def test_hlo_diff(self, test_id, model_name, overrides): """Test HLO diff for parameterized configurations.""" local_landing_dir = os.path.join(os.path.dirname(__file__), f"hlo_diff_dump_{test_id}") @@ -138,12 +138,16 @@ def test_hlo_diff(self, test_id, config_file, overrides): try: base_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "../..")) - config_path = os.path.join(base_dir, config_file) + # Load via base.yml + model_name, the normal training path: pyconfig resolves the model + # yml from model_name and merges it onto base.yml. Loading the model yml directly as the + # top-level config skips base.yml, leaving logical_axis_rules empty. + base_config_path = os.path.join(base_dir, "src/maxtext/configs/base.yml") # Arguments for train_compile test_args = [ None, - config_path, + base_config_path, + f"model_name={model_name}", "dataset_type=synthetic", "override_model_config=true", "compile_topology_num_slices=1", diff --git a/tests/integration/smoke/train_gpu_smoke_test.py b/tests/integration/smoke/train_gpu_smoke_test.py index 1df3ef630b..40dbf8ec35 100644 --- a/tests/integration/smoke/train_gpu_smoke_test.py +++ b/tests/integration/smoke/train_gpu_smoke_test.py @@ -41,7 +41,7 @@ def setUp(self): def test_tiny_config(self): test_tmpdir = os.environ.get("TEST_TMPDIR") # pylint: disable=unused-variable train_main( - [ + [ # pyrefly: ignore[bad-argument-type] None, get_test_config_path("gpu/gpu_smoke_test.yml"), # pylint: disable=f-string-without-interpolation diff --git a/tests/integration/smoke/train_int8_smoke_test.py b/tests/integration/smoke/train_int8_smoke_test.py index 274d3a124c..f31f6b976c 100644 --- a/tests/integration/smoke/train_int8_smoke_test.py +++ b/tests/integration/smoke/train_int8_smoke_test.py @@ -40,7 +40,7 @@ def setUp(self): def test_tiny_config(self): test_tmpdir = os.environ.get("TEST_TMPDIR") # pylint: disable=unused-variable train_main( - [ + [ # pyrefly: ignore[bad-argument-type] None, get_test_config_path(), # pylint: disable=f-string-without-interpolation diff --git a/tests/integration/smoke/train_using_ragged_dot_smoke_test.py b/tests/integration/smoke/train_using_ragged_dot_smoke_test.py index 8a3b7b5bcb..24035ff6e4 100644 --- a/tests/integration/smoke/train_using_ragged_dot_smoke_test.py +++ b/tests/integration/smoke/train_using_ragged_dot_smoke_test.py @@ -19,9 +19,8 @@ from absl.testing import absltest from absl.testing import parameterized - -from tests.utils.test_helpers import get_test_config_path from maxtext.trainers.pre_train import train +from tests.utils.test_helpers import get_test_config_path train_main = train.main gettempdir = tempfile.gettempdir @@ -37,41 +36,43 @@ class Train(parameterized.TestCase): def test_tiny_config(self, quantization: str): test_tmpdir = os.environ.get("TEST_TMPDIR", gettempdir()) outputs_dir = os.environ.get("TEST_UNDECLARED_OUTPUTS_DIR", test_tmpdir) - train_main([ - None, - get_test_config_path(), - f"base_output_directory={test_tmpdir}", - "run_name=ragged_dot_smoke_test", - "base_emb_dim=128", - "base_num_query_heads=4", - "base_num_kv_heads=4", - "base_mlp_dim=128", - "base_moe_mlp_dim=128", - "base_num_decoder_layers=8", - "head_dim=128", - # TODO(b/441100085): When changing the decoder_block we might - # need to adjust the tiling. - "decoder_block=deepseek", - "attention_type=mla", - "num_experts=2", - "shared_experts=1", - # Enable sparse_matmul. - "sparse_matmul=True", - # Enable ragged_dot. - "megablox=False", - f'quantization="{quantization}"', - "use_qwix_quantization=True", - "sharding_tolerance=0.15", - "per_device_batch_size=2", - "max_target_length=1024", - "dataset_type=synthetic", - "steps=10", - "enable_checkpointing=False", - "enable_goodput_recording=False", - "enable_checkpoint_cloud_logger=False", - "monitor_goodput=False", - f"metrics_file={os.path.join(outputs_dir, 'metrics.json')}", - ]) + train_main( # pyrefly: ignore[bad-argument-type] + [ + None, + get_test_config_path(), + f"base_output_directory={test_tmpdir}", + "run_name=ragged_dot_smoke_test", + "base_emb_dim=128", + "base_num_query_heads=4", + "base_num_kv_heads=4", + "base_mlp_dim=128", + "base_moe_mlp_dim=128", + "base_num_decoder_layers=8", + "head_dim=128", + # TODO(b/441100085): When changing the decoder_block we might + # need to adjust the tiling. + "decoder_block=deepseek", + "attention_type=mla", + "num_experts=2", + "shared_experts=1", + # Enable sparse_matmul. + "sparse_matmul=True", + # Enable ragged_dot. + "megablox=False", + f'quantization="{quantization}"', + "use_qwix_quantization=True", + "sharding_tolerance=0.15", + "per_device_batch_size=2", + "max_target_length=1024", + "dataset_type=synthetic", + "steps=10", + "enable_checkpointing=False", + "enable_goodput_recording=False", + "enable_checkpoint_cloud_logger=False", + "monitor_goodput=False", + f"metrics_file={os.path.join(outputs_dir, 'metrics.json')}", + ] + ) if __name__ == "__main__": diff --git a/tests/integration/sparsity_test.py b/tests/integration/sparsity_test.py index 10dd93d7c4..7154ed00f8 100644 --- a/tests/integration/sparsity_test.py +++ b/tests/integration/sparsity_test.py @@ -91,7 +91,7 @@ def test_different_quant_sparsity_configs(self, quantization: str, use_sparsity: "weight_sparsity_update_step=1", ] ) - train_main(args) + train_main(args) # pyrefly: ignore[bad-argument-type] if __name__ == "__main__": diff --git a/tests/integration/train_tests.py b/tests/integration/train_tests.py index 0fad677333..b13547125a 100644 --- a/tests/integration/train_tests.py +++ b/tests/integration/train_tests.py @@ -30,10 +30,8 @@ ) -def _small_model_base_emb_dim(decoupled, device_count): - """Return a tiny embedding dim divisible by local decoupled devices.""" - if not decoupled: - return 28 +def _small_model_base_emb_dim(device_count): + """Return a tiny embedding dim divisible by local devices.""" return ((28 + device_count - 1) // device_count) * device_count @@ -46,7 +44,7 @@ class TrainTests(unittest.TestCase): dataset_path = get_test_dataset_path() _small_model_overrides = [ - f"base_emb_dim={_small_model_base_emb_dim(decoupled, dev_count)}", + f"base_emb_dim={_small_model_base_emb_dim(dev_count)}", "base_num_query_heads=4", "base_num_kv_heads=4", "base_mlp_dim=32", @@ -98,18 +96,6 @@ class TrainTests(unittest.TestCase): rf"tokenizer_path={os.path.join(MAXTEXT_ASSETS_ROOT, 'tokenizers', 'tokenizer.llama2')}", ] + _small_model_overrides, - "tp_transpose": [ # tests base config with ici_tensor_transpose_parallelism=4 - None, - get_test_config_path(), - f"base_output_directory={_base_output_directory}", - "run_name=runner_test", - "dataset_type=synthetic", # use synthetic dataset_type to decrease training time - "steps=2", - "ici_tensor_transpose_parallelism=4", - "enable_goodput_recording=False", - rf"tokenizer_path={os.path.join(MAXTEXT_ASSETS_ROOT, 'tokenizers', 'tokenizer.llama2')}", - ] - + _small_model_overrides, "int8": [ # tests base config with int8 None, get_test_config_path(), @@ -293,6 +279,22 @@ def test_gpu_nanoo_fp8(self): def test_gpu_te_fp8_delayedscaling(self): train_main(TrainTests.CONFIGS["te_fp8_delayedscaling"] + ["attention=dot_product"]) + @pytest.mark.skip(reason="No runner with GPU arch >= 89 is available") + @pytest.mark.integration_test + @pytest.mark.gpu_only + def test_gpu_te_fp8_delayedscaling_tsp_cgemm(self): + if jax.process_count() <= 1: + pytest.skip("Requires rank-per-GPU launch (JAX_PROCESS_COUNT > 1)") + if jax.local_device_count() != 1: + pytest.skip(f"Requires rank-per-GPU launch (local_device_count==1), " f"got {jax.local_device_count()}") + if not jax.distributed.is_initialized(): + pytest.skip("Requires jax.distributed.initialize() (hardware=gpu_multiprocess)") + + train_main( + TrainTests.CONFIGS["te_fp8_delayedscaling"] + + ["attention=dot_product", "ici_tensor_sequence_parallelism=2", "use_te_comm_gemm_overlap=true"] + ) + @pytest.mark.skip(reason="No runner with GPU arch >= 89 is available") @pytest.mark.integration_test @pytest.mark.gpu_only diff --git a/tests/post_training/unit/lora_utils_test.py b/tests/post_training/unit/lora_utils_test.py index b0f229875d..e910995542 100644 --- a/tests/post_training/unit/lora_utils_test.py +++ b/tests/post_training/unit/lora_utils_test.py @@ -15,9 +15,13 @@ """Tests for Qwix LoRA utils in lora_utils.py""" import re import sys +import tempfile import unittest from unittest import mock + +from etils import epath import jax +import jax.numpy as jnp import optax import pytest from flax import nnx @@ -26,6 +30,7 @@ pytestmark = [pytest.mark.post_training] # Now safe to do top-level imports +from maxtext.common import checkpointing from tunix.sft import peft_trainer from maxtext.utils import lora_utils from maxtext.utils import model_creation_utils @@ -59,11 +64,12 @@ def _make_config(**overrides): """Return a MaxTextConfig object suitable for unit tests.""" + config_dict = _BASE_CONFIG.copy() + config_dict.update(overrides) # Use initialize_pydantic to get nested models as objects (attribute access) return pyconfig.initialize_pydantic( [sys.argv[0], get_test_config_path()], - **_BASE_CONFIG, - **overrides, + **config_dict, ) @@ -121,7 +127,12 @@ def test_build_lora_provider(self): with mock.patch("qwix.LoraProvider") as mock_provider: lora_utils._build_lora_provider(mock_config) mock_provider.assert_called_once_with( - module_path="custom/path", rank=8, alpha=16.0, dropout=0.0, weight_qtype="int8", tile_size=32 + module_path="custom/path", + rank=8, + alpha=16.0, + dropout=0.0, + weight_qtype="int8", + tile_size=32, ) def test_prepare_dummy_inputs(self): @@ -130,26 +141,41 @@ def test_prepare_dummy_inputs(self): self.assertEqual(tokens.shape, (1, 1)) self.assertEqual(positions.shape, (1, 1)) - def test_verify_lora_parameters_enabled(self): - """Test verification of LoRA parameters when enabled.""" + def test_verify_lora_parameters_success(self): + """Test verification of LoRA parameters with matches and enabled LoRA.""" mock_model = mock.MagicMock() mock_config = mock.MagicMock(spec=pyconfig.HyperParameters) + mock_config.lora = mock.MagicMock() + mock_config.lora.lora_module_path = ".*mlp/wi_0.*" + + mock_param = nnx.LoRAParam(0.0) + mock_graph_entries = [ + (("decoder", "layers", 0, "mlp", "wi_0", "lora_a"), mock_param), + ] - # Note: we use our local is_lora_enabled now - with mock.patch("maxtext.utils.lora_utils.is_lora_enabled", return_value=True): - # Should not raise + with ( + mock.patch("maxtext.utils.lora_utils.nnx.iter_graph", return_value=mock_graph_entries), + mock.patch("maxtext.utils.lora_utils.is_lora_enabled", return_value=True), + mock.patch("maxtext.utils.max_logging.log") as mock_log, + ): lora_utils._verify_lora_parameters(mock_model, mock_config) + # Should log the successful match pattern summary + log_calls = [call[0][0] for call in mock_log.call_args_list] + self.assertTrue(any("successfully matched" in msg for msg in log_calls)) + self.assertTrue(any("Sample matched submodules" in msg for msg in log_calls)) + def test_verify_lora_parameters_not_enabled_no_match(self): - """Test verification fails when LoRA parameters are expected but not found.""" + """Test verification fails with ValueError when no modules match at all.""" mock_model = mock.MagicMock() mock_config = mock.MagicMock(spec=pyconfig.HyperParameters) mock_config.lora = mock.MagicMock() - mock_config.model_name = "llama" mock_config.lora.lora_module_path = "non_existent" - with mock.patch("maxtext.utils.lora_utils.is_lora_enabled", return_value=False): - mock_model.iter_modules.return_value = [] + with ( + mock.patch("flax.nnx.iter_modules", return_value=[]), + mock.patch("maxtext.utils.lora_utils.is_lora_enabled", return_value=False), + ): with self.assertRaisesRegex(ValueError, "no LoRA parameters found"): lora_utils._verify_lora_parameters(mock_model, mock_config) @@ -173,7 +199,13 @@ def test_apply_lora_to_model_adapters_loaded(self): # If we skip Qwix, it should stay False. self.assertFalse(lora_utils.is_lora_enabled(result)) - def _run_apply_lora_test(self, scan_layers: bool, weight_qtype=None, tile_size=None, mock_multihost: bool = False): + def _run_apply_lora_test( + self, + scan_layers: bool, + weight_qtype=None, + tile_size=None, + mock_multihost: bool = False, + ): """Helper to run LoRA application test with/without scanned layers and optional QLoRA.""" # Passing nested dict as 'lora' kwarg to _make_config cfg = _make_config( @@ -246,21 +278,22 @@ def test_apply_lora_multihost_mock(self): def test_restore_lora_from_path(self): """Test restoration of LoRA parameters from a path.""" cfg = _make_config( - lora={"enable_lora": True, "lora_restore_path": "some/path", "lora_rank": 4, "lora_alpha": 8.0}, + lora={ + "enable_lora": True, + "lora_restore_path": "some/path", + "lora_rank": 4, + "lora_alpha": 8.0, + }, scan_layers=False, ) model, _ = model_creation_utils.from_pretrained(cfg, mesh=None, model_mode=model_creation_utils.MODEL_MODE_TRAIN) model = lora_utils.apply_lora_to_model(model, None, cfg) - trainer = mock.MagicMock() - trainer.model = model - trainer.train_steps = 0 - restored_state = nnx.state(model, nnx.LoRAParam) with mock.patch("orbax.checkpoint.PyTreeCheckpointer.restore", return_value=restored_state) as mock_restore: with mock.patch("flax.nnx.update") as mock_update: - lora_utils.restore_lora_from_path(trainer, cfg) + lora_utils.restore_lora_from_path(model, cfg) mock_restore.assert_called_once() args, kwargs = mock_restore.call_args self.assertEqual(args[0], "some/path") @@ -271,6 +304,135 @@ def test_restore_lora_from_path(self): self.assertTrue(kwargs["args"].partial_restore) mock_update.assert_called_once() + def test_sync_lora_metadata_default_syncs(self): + """Test that default lora rank/alpha are successfully synced from checkpoint metadata.""" + cfg = _make_config( + lora={ + "enable_lora": True, + "lora_restore_path": "dummy/path", + "lora_rank": 0, + "lora_alpha": 0.0, + } + ) + mock_metadata = mock.MagicMock() + mock_metadata.custom_metadata = {"lora": {"lora_rank": 32, "lora_alpha": 64.0}} + + with mock.patch("orbax.checkpoint.StandardCheckpointer.metadata", return_value=mock_metadata): + lora_utils.sync_lora_metadata(cfg) + self.assertEqual(cfg.lora.lora_rank, 32) + self.assertEqual(cfg.lora.lora_alpha, 64.0) + + def test_sync_lora_metadata_matching_passes(self): + """Test that matching non-default parameters pass without errors.""" + cfg = _make_config( + lora={ + "enable_lora": True, + "lora_restore_path": "dummy/path", + "lora_rank": 32, + "lora_alpha": 64.0, + } + ) + mock_metadata = mock.MagicMock() + mock_metadata.custom_metadata = {"lora": {"lora_rank": 32, "lora_alpha": 64.0}} + + with mock.patch("orbax.checkpoint.StandardCheckpointer.metadata", return_value=mock_metadata): + # Should not raise ValueError + lora_utils.sync_lora_metadata(cfg) + self.assertEqual(cfg.lora.lora_rank, 32) + self.assertEqual(cfg.lora.lora_alpha, 64.0) + + def test_sync_lora_metadata_rank_mismatch_fails(self): + """Test that configured rank mismatching checkpoint metadata rank raises ValueError.""" + cfg = _make_config( + lora={ + "enable_lora": True, + "lora_restore_path": "dummy/path", + "lora_rank": 8, + "lora_alpha": 64.0, + } + ) + mock_metadata = mock.MagicMock() + mock_metadata.custom_metadata = {"lora": {"lora_rank": 32, "lora_alpha": 64.0}} + + with mock.patch("orbax.checkpoint.StandardCheckpointer.metadata", return_value=mock_metadata): + with self.assertRaisesRegex(ValueError, "Configured lora_rank .* does not match"): + lora_utils.sync_lora_metadata(cfg) + + def test_sync_lora_metadata_alpha_mismatch_fails(self): + """Test that configured alpha mismatching checkpoint metadata alpha raises ValueError.""" + cfg = _make_config( + lora={ + "enable_lora": True, + "lora_restore_path": "dummy/path", + "lora_rank": 32, + "lora_alpha": 16.0, + } + ) + mock_metadata = mock.MagicMock() + mock_metadata.custom_metadata = {"lora": {"lora_rank": 32, "lora_alpha": 64.0}} + + with mock.patch("orbax.checkpoint.StandardCheckpointer.metadata", return_value=mock_metadata): + with self.assertRaisesRegex(ValueError, "Configured lora_alpha .* does not match"): + lora_utils.sync_lora_metadata(cfg) + + def test_save_checkpoint_passes_metadata(self): + """Test that save_checkpoint correctly generates and passes custom lora metadata to CheckpointManager.""" + cfg = _make_config( + lora={"enable_lora": True, "lora_rank": 8, "lora_alpha": 16.0}, + enable_checkpointing=True, + ) + mock_manager = mock.MagicMock() + mock_state = mock.MagicMock() + + with mock.patch("jax.block_until_ready"): + checkpointing.save_checkpoint(mock_manager, step=10, state=mock_state, config=cfg) + mock_manager.save.assert_called_once() + _, kwargs = mock_manager.save.call_args + self.assertIn("custom_metadata", kwargs) + self.assertEqual(kwargs["custom_metadata"]["lora"], cfg.lora.model_dump()) + + def test_save_and_restore_metadata_integration(self): + """Integration test checking that Orbax CheckpointManager writes and reads custom LoRA metadata.""" + + cfg_save = _make_config( + lora={"enable_lora": True, "lora_rank": 8, "lora_alpha": 16.0}, + enable_checkpointing=True, + ) + + with tempfile.TemporaryDirectory() as tmpdir: + manager = checkpointing.create_orbax_checkpoint_manager( + tmpdir, + enable_checkpointing=True, + use_async=False, + save_interval_steps=1, + use_ocdbt=False, + use_zarr3=False, + ) + + # Use save_checkpoint wrapper with a simple state + dummy_state = {"weight": jnp.array([1.0, 2.0])} + checkpointing.save_checkpoint(manager, step=0, state=dummy_state, config=cfg_save) + manager.wait_until_finished() + + # Now verify that the saved checkpoint contains metadata on disk + checkpoint_dir = epath.Path(tmpdir) / "0" + self.assertTrue((checkpoint_dir / "_CHECKPOINT_METADATA").exists()) + + # Restore using sync_lora_metadata on a config with default rank/alpha + cfg_restore = _make_config( + lora={ + "enable_lora": True, + "lora_restore_path": str(checkpoint_dir), + "lora_rank": 0, + "lora_alpha": 0.0, + } + ) + lora_utils.sync_lora_metadata(cfg_restore) + + # Verify values were successfully synced back + self.assertEqual(cfg_restore.lora.lora_rank, 8) + self.assertEqual(cfg_restore.lora.lora_alpha, 16.0) + def test_gemma4_lora_path_matching(self): """Test that the Gemma4 LoRA regex correctly matches all expected parameter paths.""" mock_config = mock.MagicMock(spec=pyconfig.HyperParameters) @@ -309,10 +471,6 @@ def test_gemma4_lora_path_matching(self): "decoder/layers_remainder/layers/0/mlp/shared_experts/wi_0/kernel", "decoder/layers_remainder/layers/0/mlp/shared_experts/wi_1/kernel", "decoder/layers_remainder/layers/0/mlp/shared_experts/wo/kernel", - # No scanned_blocks/layers_remainder prefix (e.g. fallback or direct structure) - "decoder/layers/0/self_attention/query/kernel", - "decoder/layers/0/mlp/wi_0/kernel", - "decoder/layers/layers/0/mlp/shared_experts/wi_0/kernel", ] for path in matching_paths: diff --git a/tests/post_training/unit/sft_data_processing_test.py b/tests/post_training/unit/sft_data_processing_test.py index 85baca2b7c..36ac51d3ca 100644 --- a/tests/post_training/unit/sft_data_processing_test.py +++ b/tests/post_training/unit/sft_data_processing_test.py @@ -331,7 +331,9 @@ def setUpClass(cls): ] ) if exit_code != 0: - raise ValueError(f"Download tokenizer with gcloud storage cp failed with exit code: {exit_code}") + raise unittest.SkipTest( + f"Skipping SFTDataProcessingTest: Download tokenizer with gcloud storage cp failed with exit code: {exit_code}" + ) def setUp(self): super().setUp() @@ -504,7 +506,7 @@ def setUpClass(cls): ] ) if exit_code != 0: - raise ValueError("Failed to download llama tokenizer") + raise unittest.SkipTest("Skipping SFTChatTemplateLogicTest: Failed to download llama tokenizer") def setUp(self): super().setUp() diff --git a/tests/post_training/unit/train_distill_test.py b/tests/post_training/unit/train_distill_test.py index c66ce7466a..24cd1d0bc8 100644 --- a/tests/post_training/unit/train_distill_test.py +++ b/tests/post_training/unit/train_distill_test.py @@ -50,7 +50,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive", diff --git a/tests/post_training/unit/train_rl_test.py b/tests/post_training/unit/train_rl_test.py index 0396a1f944..469f89b70c 100644 --- a/tests/post_training/unit/train_rl_test.py +++ b/tests/post_training/unit/train_rl_test.py @@ -359,6 +359,7 @@ def get_filtered_data_side_effect(dataset_name, model_tokenizer, template_config data_shuffle_seed=42, max_prefill_predict_length=10, batch_size=2, + eval_batch_size=1, num_batches=2, train_fraction=1.0, num_epoch=1, @@ -422,7 +423,8 @@ def test_prepare_datasets_with_split(self, mock_load): data_template_path="maxtext/examples/chat_templates/gsm8k_rl.json", data_shuffle_seed=42, num_batches=1, - batch_size=5, + batch_size=2, + eval_batch_size=1, train_fraction=1.0, num_epoch=1, num_test_batches=1, @@ -525,6 +527,47 @@ def test_rl_train_invalid_optimizer_memory_host_offload(self, mock_setup): with self.assertRaisesRegex(ValueError, "optimizer_memory_host_offload=True is not supported"): train_rl._rl_train_impl([], {}) # pylint: disable=protected-access + @pytest.mark.cpu_only + def test_build_reward_fns_defaults_when_no_custom(self): + """With neither knob set, the built-in 3-fn stack is returned.""" + trainer_config = SimpleNamespace(reward_functions_path="", reward_functions="") + reward_fns = train_rl.build_reward_fns(trainer_config, make_reward_fn=lambda fn: fn) + self.assertEqual( + reward_fns, + [ + train_rl.utils_rl.match_format_exactly, + train_rl.utils_rl.match_format_approximately, + train_rl.utils_rl.check_numbers, + ], + ) + + @pytest.mark.cpu_only + def test_build_reward_fns_custom_replaces_builtins(self): + """When both knobs are set, the stack is the user-provided functions only.""" + trainer_config = SimpleNamespace( + reward_functions_path="/tmp/my_rewards.py", + reward_functions="reward_a, reward_b", + ) + loaded = {"reward_a": object(), "reward_b": object()} + with mock.patch.object( + train_rl.utils_rl, "load_custom_callable", side_effect=lambda path, name: loaded[name] + ) as mock_load: + reward_fns = train_rl.build_reward_fns(trainer_config, make_reward_fn=lambda fn: fn) + self.assertEqual(reward_fns, [loaded["reward_a"], loaded["reward_b"]]) + mock_load.assert_has_calls( + [ + mock.call("/tmp/my_rewards.py", "reward_a"), + mock.call("/tmp/my_rewards.py", "reward_b"), + ] + ) + + @pytest.mark.cpu_only + def test_build_reward_fns_partial_config_falls_back(self): + """If only one of the two knobs is set, the built-in stack is used.""" + trainer_config = SimpleNamespace(reward_functions_path="/tmp/my_rewards.py", reward_functions="") + reward_fns = train_rl.build_reward_fns(trainer_config, make_reward_fn=lambda fn: fn) + self.assertEqual(len(reward_fns), 3) + class TokenizerChatTemplateTest(unittest.TestCase): """Unit tests for configure_tokenizer_chat_template.""" diff --git a/tests/unit/attention_test.py b/tests/unit/attention_test.py index 45cd1c4d27..2841933aa9 100644 --- a/tests/unit/attention_test.py +++ b/tests/unit/attention_test.py @@ -48,6 +48,7 @@ ) from maxtext.layers.attentions import Attention from maxtext.layers import embeddings +from maxtext.kernels.attention import jax_flash_attention from maxtext.configs import pyconfig from maxtext.models.qwen3 import Qwen3NextGatedDeltaNet import numpy as np @@ -57,6 +58,41 @@ from tests.utils.test_helpers import get_test_config_path +class JaxFlashAttentionTest(unittest.TestCase): + """Tests for JAX flash attention.""" + + def test_flash_attention_block_masked_soft_cap(self): + cap = 1.0 + mask_value = -1.0e9 + query = jnp.array([[[[10.0], [10.0]]]], dtype=jnp.float32) + key = jnp.array([[[[10.0], [0.0]]]], dtype=jnp.float32) + value = jnp.array([[[[1.0], [3.0]]]], dtype=jnp.float32) + mask = jnp.array([[True, True], [True, False]]) + + output = jax_flash_attention.flash_attention_block_masked( + query, + key, + value, + segment_ids=None, + block_kv=2, + block_q=1, + mask=mask, + mask_value=mask_value, + cap=cap, + ) + + logits = jnp.einsum("bhqd,bhkd->bhqk", query, key) + logits = jnp.tanh(logits / cap) * cap + logits = jnp.where(mask[None, None, :, :], logits, mask_value) + expected = jnp.einsum("bhqk,bhkd->bhqd", jax.nn.softmax(logits, axis=-1), value) + np.testing.assert_allclose( + np.asarray(output), + np.asarray(expected), + rtol=1e-6, + atol=1e-6, + ) + + class SplashLocalMaskTest(unittest.TestCase): """Tests for Splash local masks.""" @@ -292,11 +328,126 @@ def test_value_error_on_zero_chunk_size(self): _generate_chunk_attention_mask(mask_shape=(4, 4), chunk_size=0) +class LoadBalancedMaskTest(unittest.TestCase): + """Tests for load-balanced Splash masks.""" + + def test_load_balanced_local_window(self): + seq_len = 8 + window_size = 3 + q_sequence = np.asarray([0, 1, 6, 7, 2, 3, 4, 5]) + kv_sequence = np.arange(seq_len) + causal_mask = attention_op.LoadBalancedCausalMask(shape=(seq_len, seq_len), cp_size=2) + local_mask = attention_op.LoadBalancedLocalMask( + shape=(seq_len, seq_len), + window_size=(window_size - 1, window_size), + offset=0, + cp_size=2, + ) + + expected_mask = (kv_sequence[None, :] <= q_sequence[:, None]) & ( + kv_sequence[None, :] > q_sequence[:, None] - window_size + ) + + np.testing.assert_array_equal((causal_mask & local_mask)[:, :], expected_mask) + + def test_load_balanced_chunk_window(self): + seq_len = 8 + chunk_size = 2 + q_sequence = np.asarray([0, 1, 6, 7, 2, 3, 4, 5]) + kv_sequence = np.arange(seq_len) + causal_mask = attention_op.LoadBalancedCausalMask(shape=(seq_len, seq_len), cp_size=2) + chunk_mask = attention_op.LoadBalancedChunkedCausalMask( + shape=(seq_len, seq_len), + chunk_size=chunk_size, + cp_size=2, + ) + + expected_mask = (kv_sequence[None, :] <= q_sequence[:, None]) & ( + q_sequence[:, None] // chunk_size == kv_sequence[None, :] // chunk_size + ) + + np.testing.assert_array_equal((causal_mask & chunk_mask)[:, :], expected_mask) + + def test_dot_product_local_mask_uses_segment_positions(self): + config = types.SimpleNamespace(context_parallel_load_balance=True, context_sharding="context") + mesh = types.SimpleNamespace(shape={"context": 4}) + seq_len = 16 + sliding_window_size = 4 + positions = jnp.asarray(attention_op.LoadBalancedCausalMask(shape=(seq_len, seq_len), cp_size=4).q_sequence[None, :]) + query = jnp.zeros((1, seq_len, 1, 128)) + key = jnp.zeros((1, seq_len, 1, 128)) + decoder_segment_ids = jnp.ones((1, seq_len), dtype=jnp.int32) + op = AttentionOp( + config=config, + num_query_heads=1, + num_kv_heads=1, + max_target_length=seq_len, + mesh=mesh, + attention_kernel="dot_product", + attention_type=AttentionType.LOCAL_SLIDING, + sliding_window_size=sliding_window_size, + ) + + mask = op.generate_attention_mask( + query, + key, + decoder_segment_ids, + MODEL_MODE_TRAIN, + segment_positions=positions, + ) + + expected_mask = np.zeros((seq_len, seq_len), dtype=np.bool_) + for r, q_pos in enumerate(np.asarray(positions[0])): + for c, kv_pos in enumerate(np.asarray(positions[0])): + if q_pos - sliding_window_size < kv_pos <= q_pos: + expected_mask[r, c] = True + + np.testing.assert_array_equal(np.asarray(mask == 0.0)[0, 0, 0], expected_mask) + + def test_dot_product_chunk_mask_uses_segment_positions(self): + config = types.SimpleNamespace(context_parallel_load_balance=True, context_sharding="context") + mesh = types.SimpleNamespace(shape={"context": 4}) + seq_len = 16 + chunk_size = 4 + positions = jnp.asarray(attention_op.LoadBalancedCausalMask(shape=(seq_len, seq_len), cp_size=4).q_sequence[None, :]) + query = jnp.zeros((1, seq_len, 1, 128)) + key = jnp.zeros((1, seq_len, 1, 128)) + decoder_segment_ids = jnp.ones((1, seq_len), dtype=jnp.int32) + op = AttentionOp( + config=config, + num_query_heads=1, + num_kv_heads=1, + max_target_length=seq_len, + mesh=mesh, + attention_kernel="dot_product", + attention_type=AttentionType.CHUNK, + chunk_attn_window_size=chunk_size, + ) + + mask = op.generate_attention_mask( + query, + key, + decoder_segment_ids, + MODEL_MODE_TRAIN, + segment_positions=positions, + ) + + expected_mask = np.zeros((seq_len, seq_len), dtype=np.bool_) + for r, q_pos in enumerate(np.asarray(positions[0])): + for c, kv_pos in enumerate(np.asarray(positions[0])): + if q_pos >= kv_pos and q_pos // chunk_size == kv_pos // chunk_size: + expected_mask[r, c] = True + + np.testing.assert_array_equal(np.asarray(mask == 0.0)[0, 0, 0], expected_mask) + + class CudnnTePackedSequenceDescriptorTest(unittest.TestCase): """Tests packed Transformer Engine attention metadata handling.""" - def _call_packed_attention(self, sequence_descriptor): - """Runs packed TE attention with fake Transformer Engine modules.""" + def _call_te_attention( + self, sequence_descriptor, config=None, mesh=None, attention_type=AttentionType.GLOBAL, chunk_attn_window_size=None + ): + """Runs TE attention with fake Transformer Engine modules.""" sequence_descriptor.calls = [] class FakeWrappedAttention: @@ -323,17 +474,19 @@ def fake_to_nnx(*args, **kwargs): # pylint: disable=unused-argument "transformer_engine.jax.attention": attention_module, } - config = types.SimpleNamespace( - context_sharding="context", - context_parallel_strategy="ring", - context_parallel_load_balance=False, - packing=True, - dataset_type="grain", - max_segments_per_seq=4, - head_dim=2, - attention_kernel="cudnn_flash_te", - ) - mesh = types.SimpleNamespace(shape={"context": 1}) + if config is None: + config = types.SimpleNamespace( + context_sharding="context", + context_parallel_strategy="ring", + context_parallel_load_balance=False, + packing=True, + dataset_type="grain", + max_segments_per_seq=4, + head_dim=2, + attention_kernel="cudnn_flash_te", + ) + if mesh is None: + mesh = types.SimpleNamespace(shape={"context": 1}) attention = AttentionOp( config=config, mesh=mesh, @@ -342,6 +495,8 @@ def fake_to_nnx(*args, **kwargs): # pylint: disable=unused-argument num_query_heads=2, num_kv_heads=2, dtype=jnp.float32, + attention_type=attention_type, + chunk_attn_window_size=chunk_attn_window_size, ) query = jnp.zeros((1, 4, 2, 2), dtype=jnp.float32) key = jnp.zeros((1, 4, 2, 2), dtype=jnp.float32) @@ -374,7 +529,7 @@ def from_segment_ids_and_pos(cls, **kwargs): raise TypeError("older Transformer Engine does not accept THD metadata") return kwargs - output, descriptor_calls = self._call_packed_attention(SequenceDescriptor) + output, descriptor_calls = self._call_te_attention(SequenceDescriptor) self.assertEqual(len(descriptor_calls), 2) for call in descriptor_calls: @@ -383,7 +538,7 @@ def from_segment_ids_and_pos(cls, **kwargs): self.assertIs(output, descriptor_calls[0]) SequenceDescriptor.reject_thd_kwargs = True - output, descriptor_calls = self._call_packed_attention(SequenceDescriptor) + output, descriptor_calls = self._call_te_attention(SequenceDescriptor) self.assertEqual(len(descriptor_calls), 4) self.assertIn("is_thd", descriptor_calls[0]) self.assertNotIn("is_thd", descriptor_calls[1]) @@ -391,6 +546,30 @@ def from_segment_ids_and_pos(cls, **kwargs): self.assertNotIn("is_thd", descriptor_calls[3]) self.assertIs(output, descriptor_calls[1]) + def test_context_parallel_chunk_attention_rejected(self): + class SequenceDescriptor: + pass + + config = types.SimpleNamespace( + context_sharding="context", + context_parallel_strategy="all_gather", + context_parallel_load_balance=False, + packing=False, + dataset_type="synthetic", + max_segments_per_seq=1, + head_dim=2, + attention_kernel="cudnn_flash_te", + ) + mesh = types.SimpleNamespace(shape={"context": 2}) + with self.assertRaisesRegex(ValueError, "Chunk attention"): + self._call_te_attention( + SequenceDescriptor, + config=config, + mesh=mesh, + attention_type=AttentionType.CHUNK, + chunk_attn_window_size=2, + ) + class AttentionTest(parameterized.TestCase): """Test for the Attention""" diff --git a/tests/unit/checkpointing_nnx_load_test.py b/tests/unit/checkpointing_nnx_load_test.py index 69a64ec1eb..04aa898bde 100644 --- a/tests/unit/checkpointing_nnx_load_test.py +++ b/tests/unit/checkpointing_nnx_load_test.py @@ -46,6 +46,56 @@ def _abstract_nnx_state(): class TestLoadStateIfPossibleNNX(unittest.TestCase): """Cover the NNX branches in load_state_if_possible.""" + def test_emergency_linen_restore_converts_back_to_nnx(self): + kernel_abstract = jax.ShapeDtypeStruct((2, 1), jnp.float32) + step_abstract = jax.ShapeDtypeStruct((), jnp.uint32) + abstract_nnx_pure = { + "model": { + "linear": {"kernel": kernel_abstract}, + "dropout": {"rngs": {"params": {"count": step_abstract}}}, + }, + "optimizer": {"step": step_abstract}, + } + abstract_nnx_state = mock.Mock() + abstract_nnx_state.to_pure_dict.return_value = abstract_nnx_pure + restored_linen = { + "params": {"params": {"linear": {"kernel": jnp.ones((2, 1))}}}, + "step": jnp.asarray(7, dtype=jnp.int32), + } + checkpoint_manager = mock.Mock() + checkpoint_manager.restore.return_value = mock.Mock(state=restored_linen) + + restored = checkpointing._restore_emergency_linen_checkpoint_into_nnx( # pylint: disable=protected-access + checkpoint_manager, + 14, + abstract_nnx_state, + lambda leaf: ocp.type_handlers.ArrayRestoreArgs( + global_shape=leaf.shape, + dtype=leaf.dtype, + ), + ) + + checkpoint_manager.restore.assert_called_once() + restore_args = checkpoint_manager.restore.call_args.kwargs["args"].state + self.assertEqual( + set(restore_args.item.keys()), + {"params", "step"}, + ) + self.assertNotIn("model", restore_args.item) + self.assertNotIn("optimizer", restore_args.item) + self.assertTrue(restore_args.partial_restore) + self.assertIn("model", restored) + self.assertIn("optimizer", restored) + self.assertNotIn("params", restored) + self.assertNotIn("opt_state", restored) + self.assertTrue(bool(jnp.array_equal(restored["model"]["linear"]["kernel"], jnp.ones((2, 1))))) + self.assertEqual(restored["optimizer"]["step"].dtype, jnp.uint32) + self.assertEqual(int(restored["optimizer"]["step"]), 7) + self.assertEqual( + restored["model"]["dropout"]["rngs"]["params"]["count"].shape, + (), + ) + def test_load_parameters_from_path_splits_nnx_state_for_param_view(self): """When abstract_unboxed_pre_state is an nnx.State, the function must call nnx.split(model, nnx.Param, ...) to get the params and forward them to load_params_from_path.""" @@ -104,123 +154,6 @@ def test_no_paths_returns_none_none(self): self.assertIsNone(full) self.assertIsNone(params) - def test_load_state_if_possible_wraps_load_params_mismatch_exception(self): - """Verifies that load_state_if_possible intercepts and wraps PyTree mismatches in load_params_from_path.""" - abstract = _abstract_nnx_state() - with mock.patch.object( - checkpointing, - "load_params_from_path", - side_effect=ValueError("PyTree structure mismatch"), - ): - with self.assertRaises(ValueError) as ctx: - checkpointing.load_state_if_possible( - checkpoint_manager=None, - data_iterator=None, - load_parameters_from_path="gs://does-not-exist/params", - load_full_state_from_path="", - checkpoint_storage_concurrent_gb=8, - abstract_unboxed_pre_state=abstract, - ) - self.assertIn( - "Failed to load parameters from gs://does-not-exist/params.", - str(ctx.exception), - ) - self.assertIn( - "This is often caused by a mismatch in the 'scan_layers'" - " configuration", - str(ctx.exception), - ) - - def test_load_state_if_possible_re_raises_other_load_params_exceptions(self): - """Verifies that load_state_if_possible does not intercept other errors from load_params_from_path.""" - abstract = _abstract_nnx_state() - with mock.patch.object( - checkpointing, - "load_params_from_path", - side_effect=FileNotFoundError("no such file"), - ): - with self.assertRaises(FileNotFoundError): - checkpointing.load_state_if_possible( - checkpoint_manager=None, - data_iterator=None, - load_parameters_from_path="gs://does-not-exist/params", - load_full_state_from_path="", - checkpoint_storage_concurrent_gb=8, - abstract_unboxed_pre_state=abstract, - ) - - -class TestCheckpointMismatchHandling(unittest.TestCase): - """Unit tests for the checkpoint mismatch detection and wrapper context manager.""" - - def test_is_structural_or_shape_mismatch(self): - """Verifies that is_structural_or_shape_mismatch matches only shape/tree mismatches in ValueError/TypeError.""" - # Matches - self.assertTrue( - checkpointing.is_structural_or_shape_mismatch( - ValueError("PyTree structure mismatch") - ) - ) - self.assertTrue( - checkpointing.is_structural_or_shape_mismatch( - TypeError("shape mismatch in leaf") - ) - ) - self.assertTrue( - checkpointing.is_structural_or_shape_mismatch( - ValueError("tree paths matched 143/145") - ) - ) - self.assertTrue( - checkpointing.is_structural_or_shape_mismatch( - ValueError("invalid type shapedtypestruct") - ) - ) - - # Does not match - self.assertFalse( - checkpointing.is_structural_or_shape_mismatch( - ValueError("checkpoint directory does not exist") - ) - ) - self.assertFalse( - checkpointing.is_structural_or_shape_mismatch( - FileNotFoundError("file not found: checkpoint") - ) - ) - self.assertFalse( - checkpointing.is_structural_or_shape_mismatch( - RuntimeError("something went wrong") - ) - ) - - def test_handle_checkpoint_mismatch_intercepts_matching_exceptions(self): - """Verifies that handle_checkpoint_mismatch intercepts and wraps structural errors.""" - with self.assertRaises(ValueError) as ctx: - with checkpointing.handle_checkpoint_mismatch( - "load parameters", "gs://bucket/params" - ): - raise ValueError("PyTree structure mismatch") - - self.assertIn( - "Failed to load parameters from gs://bucket/params.", str(ctx.exception) - ) - self.assertIn( - "This is often caused by a mismatch in the 'scan_layers' configuration", - str(ctx.exception), - ) - self.assertIn( - "Original error: PyTree structure mismatch", str(ctx.exception) - ) - - def test_handle_checkpoint_mismatch_re_raises_non_matching_exceptions(self): - """Verifies that handle_checkpoint_mismatch does not intercept non-structural errors.""" - with self.assertRaises(FileNotFoundError): - with checkpointing.handle_checkpoint_mismatch( - "load parameters", "gs://bucket/params" - ): - raise FileNotFoundError("file not found: checkpoint") - class TestLoadParamsIntoNNX(unittest.TestCase): """Weight-only load (load_parameters_path) of a Linen-layout checkpoint into NNX.""" @@ -248,12 +181,8 @@ def test_linen_layout_params_restore_into_nnx_state(self): self.assertIsInstance(restored, nnx.State) pure = restored.to_pure_dict() - self.assertTrue( - jnp.array_equal(pure["linear"]["kernel"], weights["linear"]["kernel"]) - ) - self.assertTrue( - jnp.array_equal(pure["linear"]["bias"], weights["linear"]["bias"]) - ) + self.assertTrue(jnp.array_equal(pure["linear"]["kernel"], weights["linear"]["kernel"])) + self.assertTrue(jnp.array_equal(pure["linear"]["bias"], weights["linear"]["bias"])) if __name__ == "__main__": diff --git a/tests/unit/checkpointing_test.py b/tests/unit/checkpointing_test.py index 4b60e199e3..2907f4f314 100644 --- a/tests/unit/checkpointing_test.py +++ b/tests/unit/checkpointing_test.py @@ -305,5 +305,29 @@ def __init__(self): np.testing.assert_allclose(loaded_weight, dummy_weight) +class CheckpointMetadataTest(parameterized.TestCase): + """Tests for loading checkpoint custom metadata.""" + + @mock.patch.object(checkpointing.ocp, "StandardCheckpointer") + def test_load_checkpoint_metadata(self, mock_checkpointer_cls): + mock_ckptr = mock_checkpointer_cls.return_value + mock_metadata = mock.MagicMock() + mock_metadata.custom_metadata = {"lora": {"lora_rank": 8, "lora_alpha": 16.0}} + mock_ckptr.metadata.return_value = mock_metadata + + loaded_metadata = checkpointing.load_checkpoint_metadata("dummy/path") + self.assertEqual(loaded_metadata.get("lora"), {"lora_rank": 8, "lora_alpha": 16.0}) + mock_ckptr.metadata.assert_called_once() + + @mock.patch.object(checkpointing.ocp, "StandardCheckpointer") + def test_load_checkpoint_metadata_handles_exceptions(self, mock_checkpointer_cls): + mock_ckptr = mock_checkpointer_cls.return_value + mock_ckptr.metadata.side_effect = Exception("Checkpoint read error") + + loaded_metadata = checkpointing.load_checkpoint_metadata("corrupt/path") + self.assertEqual(loaded_metadata, {}) + mock_ckptr.metadata.assert_called_once() + + if __name__ == "__main__": absltest.main() diff --git a/tests/unit/configs_value_test.py b/tests/unit/configs_value_test.py index 0450e6671b..4d61b4caaf 100644 --- a/tests/unit/configs_value_test.py +++ b/tests/unit/configs_value_test.py @@ -96,6 +96,28 @@ def test_validation_error(self): with self.assertRaises(pydantic.ValidationError): pyconfig.initialize(argv) + def test_load_balanced_chunk_context_parallel_config(self): + argv = [ + "", + _BASE_CONFIG_PATH, + "run_name=test", + "steps=1", + "attention_type=chunk", + "chunk_attn_window_size=256", + "context_parallel_load_balance=True", + "ici_context_parallelism=2", + "hardware=tpu", + "packing=False", + "dataset_type=synthetic", + "skip_jax_distributed_system=True", + ] + mock_devices = [unittest.mock.MagicMock(slice_index=0) for _ in range(8)] + with unittest.mock.patch("jax.devices", return_value=mock_devices): + config = pyconfig.initialize(argv) + + self.assertEqual(config.attention_type, "chunk") + self.assertTrue(config.context_parallel_load_balance) + @unittest.mock.patch.dict(os.environ, {pyconfig.yaml_key_to_env_key("steps"): "123"}) def test_env_override(self): """Tests that environment variables override YAML values.""" diff --git a/tests/unit/elastic_utils_test.py b/tests/unit/elastic_utils_test.py index 5c344feb80..5e2231c66a 100644 --- a/tests/unit/elastic_utils_test.py +++ b/tests/unit/elastic_utils_test.py @@ -79,7 +79,7 @@ def setUp(self): elastic_utils.max_logging = self.fake_logging # Hook up pathwaysutils.elastic.manager.Manager to return our fake_manager - pathwaysutils.elastic.manager.Manager = lambda *args, **kwargs: self.fake_manager + pathwaysutils.elastic.manager.Manager = lambda *args, **kwargs: self.fake_manager # pyrefly: ignore[bad-assignment] pathwaysutils.elastic.manager.ScaleUpSignalError = ScaleUpSignalError # Reset global state for testing is no longer needed @@ -91,7 +91,9 @@ def tearDown(self): elastic_utils.gcs_utils = self.original_gcs_utils elastic_utils.max_logging = self.original_max_logging pathwaysutils.elastic.manager.Manager = self.original_manager_class - pathwaysutils.elastic.manager.ScaleUpSignalError = self.original_scale_up_signal_error + pathwaysutils.elastic.manager.ScaleUpSignalError = ( # pyrefly: ignore[bad-assignment] + self.original_scale_up_signal_error, + ) elastic_utils.elastic_manager = None elastic_utils.pending_reinit_recorder = None elastic_utils.pending_elastic_event_type = None @@ -245,6 +247,45 @@ def test_live_slice_indices(self): indices = elastic_utils.live_slice_indices(config) self.assertEqual(indices, {0, 1}) + def _base_mtc_keys(self, **overrides): + keys = { + "elastic_enabled": True, + "mtc_data_parallelism": 1, + "num_slices": 2, + } + keys.update(overrides) + return keys + + def test_single_controller_mtc_init_kwargs_uses_active_elastic_devices(self): + self.fake_pathwaysutils.is_pathways_backend_used.return_value = True + device0 = FakeDevice(slice_index=0) + device1 = FakeDevice(slice_index=1) + self.fake_jax.devices.return_value = [device0, device1] + self.fake_manager.active_slice_indices = {0} + + kwargs = elastic_utils.single_controller_mtc_init_kwargs(self._base_mtc_keys(mtc_data_parallelism=0)) + + self.assertEqual(kwargs["devices"], (device0,)) + self.assertEqual(kwargs["num_slices"], 1) + self.assertEqual(kwargs["data_parallelism"], 1) + + def test_single_controller_mtc_init_kwargs_raises_if_empty(self): + self.fake_pathwaysutils.is_pathways_backend_used.return_value = True + device0 = FakeDevice(slice_index=0) + self.fake_jax.devices.return_value = [device0] + self.fake_manager.active_slice_indices = {1} + + with self.assertRaisesRegex(ValueError, "Elastic single-controller MTC initialization found no active devices."): + elastic_utils.single_controller_mtc_init_kwargs(self._base_mtc_keys()) + + def test_single_controller_mtc_init_kwargs_non_elastic(self): + kwargs = elastic_utils.single_controller_mtc_init_kwargs( + self._base_mtc_keys(elastic_enabled=False, mtc_data_parallelism=3, num_slices=4) + ) + + self.assertEqual(kwargs, {"data_parallelism": 3, "num_slices": 4}) + self.assertIsNone(elastic_utils.elastic_manager) + def test_get_devices_per_host(self): device0 = FakeDevice(slice_index=0, process_index=0, task_id=0) device1 = FakeDevice(slice_index=0, process_index=0, task_id=0) @@ -355,7 +396,7 @@ def test_record_elastic_wait_end_and_reinit_start_noop_on_first_attempt(self): def test_record_elastic_wait_end_and_reinit_start(self): """Test recording end of slice down and start of reinit.""" - elastic_utils.pending_elastic_event_type = "elastic_slice_down" + elastic_utils.pending_elastic_event_type = "elastic_slice_down" # pyrefly: ignore[bad-assignment] fake_recorder = Mock() elastic_utils.record_elastic_wait_end_and_reinit_start(fake_recorder) diff --git a/tests/unit/grain_data_processing_test.py b/tests/unit/grain_data_processing_test.py index dc129b5ebc..d061f5b10d 100644 --- a/tests/unit/grain_data_processing_test.py +++ b/tests/unit/grain_data_processing_test.py @@ -181,6 +181,7 @@ def _make_config(self, **overrides): return pyconfig.initialize([sys.argv[0], get_test_config_path()], **kwargs) +@pytest.mark.cpu_only class GrainArrayRecordProcessingTest( _GrainArrayRecordSetup, GrainDeterminismMixin, GrainBaseProcessingTest, unittest.TestCase ): @@ -201,6 +202,7 @@ def test_batch_determinism(self): super().test_batch_determinism() +@pytest.mark.cpu_only class GrainArrayRecordProcessingWithMultiSourceBlendingTest( _GrainArrayRecordSetup, GrainBaseProcessingTest, unittest.TestCase ): @@ -211,6 +213,7 @@ def setUp(self): self.config = self._make_config(grain_train_files=train_files_weighted) +@pytest.mark.cpu_only class GrainArrayRecordProcessingWithMixtureConfigTest(_GrainArrayRecordSetup, GrainBaseProcessingTest, unittest.TestCase): def setUp(self): @@ -270,6 +273,7 @@ def setUp(self): self.config = self._make_config(grain_ram_budget_mb=512, grain_worker_count=-1) # Enable auto-tuning +@pytest.mark.cpu_only class GrainArrayRecordTiktokenTest(_GrainArrayRecordSetup, GrainBaseProcessingTest, unittest.TestCase): """Test grain data processing with tiktoken tokenizer.""" @@ -281,6 +285,7 @@ def setUp(self): ) +@pytest.mark.cpu_only class GrainArrayRecordHFTokenizerTest(_GrainArrayRecordSetup, GrainBaseProcessingTest, unittest.TestCase): """Test grain data processing with HuggingFace tokenizer.""" @@ -292,6 +297,7 @@ def setUp(self): ) +@pytest.mark.cpu_only class GrainArrayRecordBestFitPackingTest(_GrainArrayRecordSetup, GrainBaseProcessingTest, unittest.TestCase): """Test grain data processing with best_fit packing strategy.""" @@ -381,6 +387,7 @@ def _make_config(self, **overrides): return pyconfig.initialize([sys.argv[0], get_test_config_path()], **kwargs) +@pytest.mark.cpu_only class GrainParquetProcessingTest(_GrainParquetSetup, GrainDeterminismMixin, GrainBaseProcessingTest, unittest.TestCase): """Test grain data processing with Parquet format. @@ -519,6 +526,7 @@ def _make_config(self, **overrides): return pyconfig.initialize([sys.argv[0], get_test_config_path()], **kwargs) +@pytest.mark.cpu_only class GrainTFRecordProcessingTest(_GrainTFRecordSetup, GrainDeterminismMixin, GrainBaseProcessingTest, unittest.TestCase): """Test grain data processing with TFRecord format. @@ -533,6 +541,7 @@ def setUpClass(cls): super().setUpClass() +@pytest.mark.cpu_only class GrainTFRecordPreTokenizedProcessingTest(_GrainTFRecordSetup, GrainBaseProcessingTest, unittest.TestCase): """Test grain data processing with a pre-tokenized TFRecord dataset (tokenize_train_data=False). @@ -552,5 +561,29 @@ def setUp(self): ) +@pytest.mark.cpu_only +class GrainFewerFilesThanHostsTest(_GrainTFRecordSetup, GrainBaseProcessingTest, unittest.TestCase): + """Tests data loading when file count < dataloading_host_count. + + _GrainTFRecordSetup provides a single TFRecord file. Overriding process_indices + to [0, 1] makes make_grain_train_iterator use dataloading_host_count=2, simulating + the undersized scenario without a real multi-host runner. The inherited test_train_ds + then validates the full batch shape through this code path. + """ + + def setUp(self): + super().setUp() + # Simulate 2 dataloading hosts with only 1 file to trigger the + # fewer-files-than-hosts path in get_datasets via make_grain_train_iterator. + self.process_indices = [0, 1] + + def test_raises_when_grain_worker_count_exceeds_files_per_host(self): + # 1 file, 2 hosts (via process_indices=[0,1]) → files_per_host=1; + # grain_worker_count=2 must raise. + config = self._make_config(grain_worker_count=2) + with self.assertRaises(ValueError): + grain_data_processing.make_grain_train_iterator(config, self.mesh, self.process_indices) + + if __name__ == "__main__": unittest.main() diff --git a/tests/unit/hf_checkpoint_conversion_test.py b/tests/unit/hf_checkpoint_conversion_test.py index 6451a50fd5..a920e0a762 100644 --- a/tests/unit/hf_checkpoint_conversion_test.py +++ b/tests/unit/hf_checkpoint_conversion_test.py @@ -367,6 +367,7 @@ def test_get_maxtext_model_info(self): "hidden_size_per_layer_input=128", "vocab_size_per_layer_input=256", "vocab_size=256", + "skip_jax_distributed_system=True", ], override_model_config=True, ) @@ -417,7 +418,7 @@ def test_recursive_update(self): @unittest.mock.patch("maxtext.checkpoint_conversion.utils.utils.ocp.Checkpointer") @unittest.mock.patch("maxtext.checkpoint_conversion.utils.utils.epath.Path") @unittest.mock.patch("maxtext.checkpoint_conversion.utils.utils.jax.devices") - def test_load_orbax_checkpoint_recursive_merge(self, mock_jax_devices, mock_path, mock_checkpointer_cls): + def test_load_orbax_checkpoint_recursive_merge(self, mock_jax_devices, _mock_path, mock_checkpointer_cls): # Mock jax devices mock_jax_devices.return_value = [MagicMock()] diff --git a/tests/unit/hf_data_processing_test.py b/tests/unit/hf_data_processing_test.py index 262c56ff9b..93564514a3 100644 --- a/tests/unit/hf_data_processing_test.py +++ b/tests/unit/hf_data_processing_test.py @@ -19,6 +19,7 @@ import os.path import jax +import pytest from jax.sharding import Mesh from jax.experimental import mesh_utils @@ -30,6 +31,7 @@ from tests.utils.test_helpers import get_test_config_path, get_test_base_output_directory +@pytest.mark.cpu_only class HfDataProcessingTest(unittest.TestCase): def setUp(self): diff --git a/tests/unit/max_utils_test.py b/tests/unit/max_utils_test.py index 41572529d6..a10af2eb95 100644 --- a/tests/unit/max_utils_test.py +++ b/tests/unit/max_utils_test.py @@ -492,6 +492,44 @@ def test_single_controller_multi_tier_checkpointing_uses_colocated_python(self, use_colocated_python=True, ) + @mock.patch("maxtext.utils.max_utils.elastic_utils.single_controller_mtc_init_kwargs") + @mock.patch("maxtext.utils.max_utils.initialize_multi_tier_checkpointing") + @mock.patch("jax.distributed.initialize") + def test_single_controller_multi_tier_checkpointing_uses_elastic_utils_kwargs( + self, mock_init, mock_mtc, mock_mtc_init_kwargs + ): + active_devices = ( + mock.Mock(slice_index=0), + mock.Mock(slice_index=0), + ) + mock_mtc_init_kwargs.return_value = { + "data_parallelism": 1, + "num_slices": 1, + "devices": active_devices, + } + raw_keys = self._base_keys( + enable_single_controller=True, + enable_multi_tier_checkpointing=True, + elastic_enabled=True, + mtc_data_parallelism=0, + num_slices=2, + ) + + max_utils.maybe_initialize_jax_distributed_system(raw_keys) + + mock_init.assert_not_called() + mock_mtc_init_kwargs.assert_called_once_with(raw_keys) + mock_mtc.assert_called_once_with( + local_checkpoint_directory=self._base_keys()["local_checkpoint_directory"], + backup_interval_minutes=self._base_keys()["multi_tier_checkpointing_backup_interval_minutes"], + run_name=self._base_keys()["run_name"], + jax_initialization_timeout_seconds=self._base_keys()["jax_distributed_initialization_timeout"], + data_parallelism=1, + num_slices=1, + use_colocated_python=True, + devices=active_devices, + ) + @mock.patch("jax.distributed.initialize") def test_tpu_checkpointing_no_emergency_calls_jax_init(self, mock_init): raw_keys = self._base_keys(enable_checkpointing=True, compile_topology_num_slices=-1) diff --git a/tests/unit/maxtext_utils_test.py b/tests/unit/maxtext_utils_test.py index 0d51e1e1c5..cce1da530d 100644 --- a/tests/unit/maxtext_utils_test.py +++ b/tests/unit/maxtext_utils_test.py @@ -352,9 +352,7 @@ def setUp(self): self.mesh = Mesh(devices_array, self.config.mesh_axes) quant = quantizations.configure_quantization(self.config) if self.config.pure_nnx: - self._create_model_partial, self.model = ( - model_creation_utils.create_nnx_abstract_model(self.config, self.mesh) - ) + self._create_model_partial, self.model = model_creation_utils.create_nnx_abstract_model(self.config, self.mesh) else: self.model = models.transformer_as_linen(self.config, mesh=self.mesh, quant=quant, model_mode=MODEL_MODE_TRAIN) @@ -389,7 +387,8 @@ def create_train_state_fn(): init_state_fn = create_train_state_fn else: init_state_fn = functools.partial(maxtext_utils.init_initial_state, self.model, tx, self.config, True, rng) - state, _, _, _ = maxtext_utils.setup_initial_state(None, self.config, self.mesh, None, init_state_fn) + state, _, _, _, was_restored = maxtext_utils.setup_initial_state(None, self.config, self.mesh, None, init_state_fn) + self.assertFalse(was_restored) if self.config.pure_nnx: self.assertIsNotNone(state.optimizer) else: @@ -951,7 +950,21 @@ class MockConfig: def setUp(self): # Setup a dummy device array for the mock to return - self.devices_array = np.array(jax.devices()) + devices = jax.devices() + if len(devices) < 2: + mock_devices = [] + for i in range(2): + d = MagicMock(spec=jax.Device) + d.id = i + d.device_kind = "cpu" + d.platform = "cpu" + # make it hashable + d.__hash__ = lambda self, idx=i: idx + d.__eq__ = lambda self, other, idx=i: isinstance(other, MagicMock) and other.id == idx + mock_devices.append(d) + self.devices_array = np.array(mock_devices) + else: + self.devices_array = np.array(devices) @patch("maxtext.utils.maxtext_utils.create_device_mesh") def test_get_mesh_explicit_mode(self, mock_create_device_mesh): @@ -1377,13 +1390,9 @@ def setUp(self): self.mesh = Mesh(devices_array, self.config.mesh_axes) quant = quantizations.configure_quantization(self.config) if self.config.pure_nnx: - self._create_model_partial, self.model = ( - model_creation_utils.create_nnx_abstract_model(self.config, self.mesh) - ) + self._create_model_partial, self.model = model_creation_utils.create_nnx_abstract_model(self.config, self.mesh) else: - self.model = Transformer( - self.config, mesh=self.mesh, quant=quant, model_mode=MODEL_MODE_TRAIN - ) + self.model = Transformer(self.config, mesh=self.mesh, quant=quant, model_mode=MODEL_MODE_TRAIN) def test_setup_training_state_returns_train_state(self): rng = jax.random.PRNGKey(0) @@ -1405,7 +1414,8 @@ def create_train_state_fn(): True, rng, ) - state, _, _, _ = maxtext_utils.setup_training_state(None, self.config, self.mesh, None, init_state_fn) + state, _, _, _, was_restored = maxtext_utils.setup_training_state(None, self.config, self.mesh, None, init_state_fn) + self.assertFalse(was_restored) if self.config.pure_nnx: self.assertIsNotNone(state.optimizer) else: @@ -1422,13 +1432,9 @@ def setUp(self): self.mesh = Mesh(devices_array, self.config.mesh_axes) quant = quantizations.configure_quantization(self.config) if self.config.pure_nnx: - self._create_model_partial, self.model = ( - model_creation_utils.create_nnx_abstract_model(self.config, self.mesh) - ) + self._create_model_partial, self.model = model_creation_utils.create_nnx_abstract_model(self.config, self.mesh) else: - self.model = Transformer( - self.config, mesh=self.mesh, quant=quant, model_mode=MODEL_MODE_TRAIN - ) + self.model = Transformer(self.config, mesh=self.mesh, quant=quant, model_mode=MODEL_MODE_TRAIN) self.rng = jax.random.PRNGKey(0) self.tx = optax.adam(learning_rate=0.001) @@ -1442,9 +1448,7 @@ def create_train_state_fn(): init_state_fn = create_train_state_fn annotations = maxtext_utils_nnx.get_partition_spec_nnx( - maxtext_utils.get_abstract_state( - self.config, self.mesh, init_state_fn, True - )[2] + maxtext_utils.get_abstract_state(self.config, self.mesh, init_state_fn, True)[2] ) else: init_state_fn = functools.partial( @@ -1455,9 +1459,7 @@ def create_train_state_fn(): True, self.rng, ) - annotations = maxtext_utils.get_logical_annotations( - self.config, self.mesh, init_state_fn - ) + annotations = maxtext_utils.get_logical_annotations(self.config, self.mesh, init_state_fn) # Result should be a pytree with PartitionSpec leaves leaves = jax.tree_util.tree_leaves(annotations) self.assertGreater(len(leaves), 0) @@ -1593,9 +1595,7 @@ class MockConfig: optimizer_memory_host_offload: bool = False parameter_memory_host_offload: bool = False param_scan_axis: int = 0 - logical_axis_rules: list = field( - default_factory=lambda: [["data", ["data"]], ["model", ["model"]]] - ) + logical_axis_rules: list = field(default_factory=lambda: [["data", ["data"]], ["model", ["model"]]]) class MockTrainState(nnx.Module): """Simulates a TrainState with params and optimizer state.""" @@ -1692,5 +1692,43 @@ def test_invalid_init_fn(self): maxtext_utils.get_abstract_state_nnx(self.config, self.mesh, None) +class TestKVCacheScanHelpers(unittest.TestCase): + """Tests for KV cache scan helper functions.""" + + def test_prepare_kv_caches_for_scan_valid(self): + kv_caches = [jnp.array([1.0]), jnp.array([2.0]), jnp.array([3.0]), jnp.array([4.0])] + # scan_length = 2, block_len = 2 + grouped = maxtext_utils.prepare_kv_caches_for_scan(kv_caches, scan_length=2, block_len=2, stack=False) + self.assertEqual(len(grouped), 2) + self.assertEqual(grouped[0], (kv_caches[0], kv_caches[1])) + self.assertEqual(grouped[1], (kv_caches[2], kv_caches[3])) + + def test_prepare_kv_caches_for_scan_invalid_type(self): + kv_caches_tuple = (jnp.array([1.0]), jnp.array([2.0])) + with self.assertRaises(TypeError): + maxtext_utils.prepare_kv_caches_for_scan(kv_caches_tuple, scan_length=1, block_len=2) + + kv_caches_array = jnp.array([1.0, 2.0]) + with self.assertRaises(TypeError): + maxtext_utils.prepare_kv_caches_for_scan(kv_caches_array, scan_length=1, block_len=2) + + def test_update_kv_caches_after_scan_valid(self): + kv_caches = [jnp.array([1.0]), jnp.array([2.0]), jnp.array([3.0]), jnp.array([4.0])] + returned_kv_cache = [(jnp.array([10.0]), jnp.array([20.0])), (jnp.array([30.0]), jnp.array([40.0]))] + + maxtext_utils.update_kv_caches_after_scan(kv_caches, returned_kv_cache, scan_length=2, block_len=2, stacked=False) + + self.assertTrue(jnp.array_equal(kv_caches[0], jnp.array([10.0]))) + self.assertTrue(jnp.array_equal(kv_caches[1], jnp.array([20.0]))) + self.assertTrue(jnp.array_equal(kv_caches[2], jnp.array([30.0]))) + self.assertTrue(jnp.array_equal(kv_caches[3], jnp.array([40.0]))) + + def test_update_kv_caches_after_scan_invalid_type(self): + kv_caches_tuple = (jnp.array([1.0]), jnp.array([2.0])) + returned_kv_cache = [(jnp.array([10.0]), jnp.array([20.0]))] + with self.assertRaises(TypeError): + maxtext_utils.update_kv_caches_after_scan(kv_caches_tuple, returned_kv_cache, scan_length=1, block_len=2) + + if __name__ == "__main__": unittest.main() diff --git a/tests/unit/model_creation_utils_test.py b/tests/unit/model_creation_utils_test.py index 2568547944..b0ec54f694 100644 --- a/tests/unit/model_creation_utils_test.py +++ b/tests/unit/model_creation_utils_test.py @@ -723,6 +723,65 @@ def test_checkpoint_load_error_propagates(self, mock_ocp): with self.assertRaises(RuntimeError): model_creation_utils.from_pretrained(cfg, self.mesh) + @patch("maxtext.utils.model_creation_utils.checkpointing.load_checkpoint_metadata") + def test_scan_layers_mismatch_raises_error(self, mock_load_meta): + """ValueError is raised if run specifies scan_layers=True but checkpoint specifies scan_layers=False.""" + mock_load_meta.return_value = {"scan_layers": False} + + cfg = _make_config( + enable_checkpointing=True, load_parameters_path="gs://fake/scan_layers_false_ckpt", scan_layers=True + ) + + with self.assertRaises(ValueError) as context: + model_creation_utils.from_pretrained(cfg, self.mesh) + self.assertIn( + "Configuration mismatch: Your run specifies scan_layers=True, " + "but the checkpoint was saved with scan_layers=False", + str(context.exception), + ) + + @patch("maxtext.utils.model_creation_utils.checkpointing.load_checkpoint_metadata") + @patch("maxtext.utils.model_creation_utils.ocp") + def test_scan_layers_match_no_error(self, mock_ocp, mock_load_meta): + """If the run specifies scan_layers=True and the checkpoint matches, it proceeds without error.""" + mock_load_meta.return_value = {"scan_layers": True} + + mock_ckptr = MagicMock() + mock_ckptr.metadata.return_value = self._make_linen_metadata_mock() + mock_ckptr.restore.side_effect = lambda path, item=None, **kw: item + mock_ocp.Checkpointer.return_value = mock_ckptr + mock_ocp.PyTreeCheckpointHandler.return_value = MagicMock() + mock_ocp.checkpoint_utils.construct_restore_args.return_value = {} + mock_ocp.ArrayRestoreArgs = ocp.ArrayRestoreArgs + + cfg = _make_config( + enable_checkpointing=True, load_parameters_path="gs://fake/scan_layers_true_ckpt", scan_layers=True + ) + + model = model_creation_utils.from_pretrained(cfg, self.mesh) + self.assertIsInstance(model, models.Transformer) + + @patch("maxtext.utils.model_creation_utils.checkpointing.load_checkpoint_metadata") + @patch("maxtext.utils.model_creation_utils.ocp") + def test_scan_layers_missing_metadata_no_error(self, mock_ocp, mock_load_meta): + """Skip verification and proceed if custom_metadata lacks 'scan_layers'.""" + mock_load_meta.return_value = {} + + mock_ckptr = MagicMock() + mock_ckptr.metadata.return_value = self._make_linen_metadata_mock() + mock_ckptr.restore.side_effect = lambda path, item=None, **kw: item + mock_ocp.Checkpointer.return_value = mock_ckptr + mock_ocp.PyTreeCheckpointHandler.return_value = MagicMock() + mock_ocp.checkpoint_utils.construct_restore_args.return_value = {} + mock_ocp.ArrayRestoreArgs = ocp.ArrayRestoreArgs + + cfg = _make_config( + enable_checkpointing=True, load_parameters_path="gs://fake/scan_layers_missing_ckpt", scan_layers=True + ) + + model = model_creation_utils.from_pretrained(cfg, self.mesh) + self.assertIsInstance(model, models.Transformer) + class TestSetupDecodeStateFromNnx(unittest.TestCase): """Tests for setup_decode_state_from_nnx().""" @@ -773,5 +832,190 @@ def test_returns_linen_train_state_and_annotations(self): self.assertIsNotNone(state_mesh_annotations) +class TestVerifyAndSyncScanLayers(unittest.TestCase): + """Tests for verify_and_sync_scan_layers().""" + + def setUp(self): + self.mesh = Mesh(np.array(jax.devices()[:1]), axis_names=("x",)) + + @patch("maxtext.utils.model_creation_utils.checkpointing.load_checkpoint_metadata") + def test_sync_to_false_when_implicit(self, mock_load_meta): + """If scan_layers is not explicit, sync scan_layers to False from checkpoint metadata.""" + mock_load_meta.return_value = {"scan_layers": False} + # Create config without scan_layers in kwargs so it's not explicit + cfg = _make_config(enable_checkpointing=True, load_parameters_path="gs://fake/ckpt") + + # Pre-assertions + pydantic_cfg = getattr(cfg, "_pydantic_config", cfg) + self.assertTrue(cfg.scan_layers) # default is True + self.assertNotIn("scan_layers", pydantic_cfg.model_fields_set) + + # Call verify_and_sync_scan_layers + synced_cfg = model_creation_utils.verify_and_sync_scan_layers(cfg) + + # Post-assertions + self.assertFalse(synced_cfg.scan_layers) + synced_pydantic_cfg = getattr(synced_cfg, "_pydantic_config", synced_cfg) + self.assertFalse(synced_pydantic_cfg.scan_layers) + + @patch("maxtext.utils.model_creation_utils.checkpointing.load_checkpoint_metadata") + def test_sync_to_true_when_implicit(self, mock_load_meta): + """If scan_layers is not explicit, sync scan_layers to True from checkpoint metadata.""" + mock_load_meta.return_value = {"scan_layers": True} + cfg = _make_config(enable_checkpointing=True, load_parameters_path="gs://fake/ckpt") + + synced_cfg = model_creation_utils.verify_and_sync_scan_layers(cfg) + self.assertTrue(synced_cfg.scan_layers) + + @patch("maxtext.utils.model_creation_utils.checkpointing.load_checkpoint_metadata") + def test_explicit_match_raises_no_error(self, mock_load_meta): + """If scan_layers is explicit and matches checkpoint metadata, no error is raised.""" + mock_load_meta.return_value = {"scan_layers": True} + cfg = _make_config(enable_checkpointing=True, load_parameters_path="gs://fake/ckpt", scan_layers=True) + + # Pre-assertions + pydantic_cfg = getattr(cfg, "_pydantic_config", cfg) + self.assertTrue(cfg.scan_layers) + self.assertIn("scan_layers", pydantic_cfg.model_fields_set) + + synced_cfg = model_creation_utils.verify_and_sync_scan_layers(cfg) + self.assertTrue(synced_cfg.scan_layers) + + @patch("maxtext.utils.model_creation_utils.checkpointing.load_checkpoint_metadata") + def test_explicit_mismatch_raises_value_error(self, mock_load_meta): + """If scan_layers is explicit and mismatches checkpoint metadata, ValueError is raised.""" + mock_load_meta.return_value = {"scan_layers": False} + cfg = _make_config(enable_checkpointing=True, load_parameters_path="gs://fake/ckpt", scan_layers=True) + + # Pre-assertions + pydantic_cfg = getattr(cfg, "_pydantic_config", cfg) + self.assertTrue(cfg.scan_layers) + self.assertIn("scan_layers", pydantic_cfg.model_fields_set) + + with self.assertRaises(ValueError) as context: + model_creation_utils.verify_and_sync_scan_layers(cfg) + + self.assertIn("Configuration mismatch", str(context.exception)) + + @patch("maxtext.utils.model_creation_utils.checkpointing.load_checkpoint_metadata") + @patch("maxtext.utils.model_creation_utils.max_logging.log") + def test_sync_log_and_early_return(self, mock_log, mock_load_meta): + """Test that we log only when auto-resolution actually changes scan_layers, and early return otherwise.""" + # Scenario A: saved_scan_layers == config.scan_layers (default True). + # Since they match, it should return early, and NO log should be printed. + mock_load_meta.return_value = {"scan_layers": True} + cfg = _make_config(enable_checkpointing=True, load_parameters_path="gs://fake/ckpt") + + synced_cfg = model_creation_utils.verify_and_sync_scan_layers(cfg) + self.assertTrue(synced_cfg.scan_layers) + mock_log.assert_not_called() + + # Scenario B: saved_scan_layers (False) != config.scan_layers (default True), and is not explicit (implicit). + # It should log the auto-resolution and update scan_layers to False. + mock_load_meta.return_value = {"scan_layers": False} + cfg2 = _make_config(enable_checkpointing=True, load_parameters_path="gs://fake/ckpt") + + synced_cfg2 = model_creation_utils.verify_and_sync_scan_layers(cfg2) + self.assertFalse(synced_cfg2.scan_layers) + mock_log.assert_called_once_with("Setting scan_layers=False loaded from checkpoint metadata.") + + if __name__ == "__main__": unittest.main() + + +class TestFromPretrainedAuth(unittest.TestCase): + """Tests for Hugging Face authentication in from_pretrained.""" + + def _make_nnx_metadata_mock(self): + meta = MagicMock() + meta.item_metadata.tree.keys.return_value = ["decoder"] + meta.item_metadata.tree.get.return_value = {} + return meta + + @patch("maxtext.utils.model_creation_utils.ocp") + @patch("maxtext.utils.model_creation_utils.subprocess.run") + @patch("maxtext.utils.model_creation_utils.get_token") + @patch("maxtext.utils.model_creation_utils.epath.Path") + def test_auth_success_with_config_token(self, mock_path, mock_get_token, mock_run, mock_ocp): + config = _make_config( + convert_checkpoint_if_possible=True, + base_output_directory="gs://fake_bucket/fake_run", + hf_access_token="config_token", + ) + mesh = _make_mesh(config) + + mock_ckpt_path = MagicMock() + mock_ckpt_path.exists.return_value = False + mock_path.return_value.__truediv__.return_value.__truediv__.return_value = mock_ckpt_path + + mock_ckptr = MagicMock() + mock_ckptr.metadata.return_value = self._make_nnx_metadata_mock() + mock_ckptr.restore.side_effect = lambda path, item=None, **kw: item + mock_ocp.Checkpointer.return_value = mock_ckptr + mock_ocp.checkpoint_utils.construct_restore_args.return_value = {} + mock_ocp.ArrayRestoreArgs = ocp.ArrayRestoreArgs + + model = model_creation_utils.from_pretrained(config, mesh) + self.assertIsInstance(model, models.Transformer) + + mock_run.assert_called_once() + called_env = mock_run.call_args[1].get("env", {}) + self.assertEqual(called_env.get("HF_TOKEN"), "config_token") + mock_get_token.assert_not_called() + + @patch("maxtext.utils.model_creation_utils.ocp") + @patch("maxtext.utils.model_creation_utils.subprocess.run") + @patch("maxtext.utils.model_creation_utils.get_token") + @patch("maxtext.utils.model_creation_utils.epath.Path") + def test_auth_success_with_cached_token(self, mock_path, mock_get_token, mock_run, mock_ocp): + config = _make_config( + convert_checkpoint_if_possible=True, + base_output_directory="gs://fake_bucket/fake_run", + hf_access_token="", + ) + mesh = _make_mesh(config) + + mock_ckpt_path = MagicMock() + mock_ckpt_path.exists.return_value = False + mock_path.return_value.__truediv__.return_value.__truediv__.return_value = mock_ckpt_path + + mock_get_token.return_value = "cached_token" + + mock_ckptr = MagicMock() + mock_ckptr.metadata.return_value = self._make_nnx_metadata_mock() + mock_ckptr.restore.side_effect = lambda path, item=None, **kw: item + mock_ocp.Checkpointer.return_value = mock_ckptr + mock_ocp.checkpoint_utils.construct_restore_args.return_value = {} + mock_ocp.ArrayRestoreArgs = ocp.ArrayRestoreArgs + + model = model_creation_utils.from_pretrained(config, mesh) + self.assertIsInstance(model, models.Transformer) + + mock_get_token.assert_called_once() + mock_run.assert_called_once() + called_env = mock_run.call_args[1].get("env", {}) + self.assertEqual(called_env.get("HF_TOKEN"), "cached_token") + + @patch("maxtext.utils.model_creation_utils.ocp") + @patch("maxtext.utils.model_creation_utils.subprocess.run") + @patch("maxtext.utils.model_creation_utils.get_token") + @patch("maxtext.utils.model_creation_utils.epath.Path") + def test_auth_failure_no_token(self, mock_path, mock_get_token, mock_run, mock_ocp): + config = _make_config( + convert_checkpoint_if_possible=True, + base_output_directory="gs://fake_bucket/fake_run", + hf_access_token="", + ) + mesh = _make_mesh(config) + + mock_ckpt_path = MagicMock() + mock_ckpt_path.exists.return_value = False + mock_path.return_value.__truediv__.return_value.__truediv__.return_value = mock_ckpt_path + + mock_get_token.return_value = None + + with self.assertRaisesRegex(ValueError, "hf_access_token must be provided"): + model_creation_utils.from_pretrained(config, mesh) + + mock_run.assert_not_called() diff --git a/tests/unit/moe_test.py b/tests/unit/moe_test.py index 4f66a045f1..c80ec56ad4 100644 --- a/tests/unit/moe_test.py +++ b/tests/unit/moe_test.py @@ -701,7 +701,7 @@ def loss_fn(params, x): # test). Without checking this, DCE removes the bwd entirely. self.assertEqual(x_grad_ref.shape, x_grad_rs.shape, "Hidden-state grad shape mismatch") self.assertTrue( - jnp.allclose(x_grad_rs.astype(jnp.float32), x_grad_ref.astype(jnp.float32), rtol=1e-2, atol=1e-2), + jnp.allclose(x_grad_rs.astype(jnp.float32), x_grad_ref.astype(jnp.float32), rtol=1e-2, atol=8e-2), msg=( "Hidden-state gradient mismatch: max abs diff=" f"{jnp.max(jnp.abs(x_grad_rs.astype(jnp.float32) - x_grad_ref.astype(jnp.float32)))}" @@ -715,7 +715,7 @@ def loss_fn(params, x): for i, (g_ref, g_rs) in enumerate(zip(leaves_ref, leaves_rs)): self.assertEqual(g_ref.shape, g_rs.shape, f"Grad shape mismatch at leaf {i}") self.assertTrue( - jnp.allclose(g_rs.astype(jnp.float32), g_ref.astype(jnp.float32), rtol=1e-2, atol=1e-2), + jnp.allclose(g_rs.astype(jnp.float32), g_ref.astype(jnp.float32), rtol=1e-2, atol=8e-2), msg=( f"Gradient mismatch at leaf {i} (shape={g_ref.shape}): " f"max abs diff={jnp.max(jnp.abs(g_rs.astype(jnp.float32) - g_ref.astype(jnp.float32)))}" @@ -723,34 +723,28 @@ def loss_fn(params, x): ) @pytest.mark.tpu_only - @pytest.mark.skip_on_tpu7x def test_ragged_sort_loss_and_grad_ring_of_experts(self): self._run_ragged_sort_loss_and_grad(use_ring_of_experts=True) @pytest.mark.tpu_only - @pytest.mark.skip_on_tpu7x def test_ragged_sort_loss_and_grad_ring_of_experts_ragged_buffer(self): self._run_ragged_sort_loss_and_grad(use_ring_of_experts=True, ragged_buffer_factor=1.5) @pytest.mark.tpu_only - @pytest.mark.skip_on_tpu7x def test_ragged_sort_loss_and_grad_ring_of_experts_fallback(self): self._run_ragged_sort_loss_and_grad( use_ring_of_experts=True, ragged_gather_fallback=True, ragged_gather_reduce_fallback=True ) @pytest.mark.tpu_only - @pytest.mark.skip_on_tpu7x def test_ragged_sort_loss_and_grad_no_ring_of_experts(self): self._run_ragged_sort_loss_and_grad(use_ring_of_experts=False) @pytest.mark.tpu_only - @pytest.mark.skip_on_tpu7x def test_ragged_sort_loss_and_grad_no_ring_of_experts_ragged_buffer(self): self._run_ragged_sort_loss_and_grad(use_ring_of_experts=False, ragged_buffer_factor=1.5) @pytest.mark.tpu_only - @pytest.mark.skip_on_tpu7x def test_ragged_sort_loss_and_grad_no_ring_of_experts_fallback(self): self._run_ragged_sort_loss_and_grad( use_ring_of_experts=False, ragged_gather_fallback=True, ragged_gather_reduce_fallback=True @@ -795,51 +789,6 @@ def test_moe_fsdp_two_stage_parallelism_tpu_only(self): actual_output, _, _ = self.get_moe_output(variables, hidden_states, cfg, mesh) self.assertTrue(jax.numpy.allclose(expected_output, actual_output, rtol=1e-02, atol=1e-02, equal_nan=False)) - @pytest.mark.tpu_only - def test_megablox_tp_transpose_parallelism(self): - cfg = pyconfig.initialize( - [None, get_test_config_path()], - run_name="moe_block_megablox_tp_transpose_test", - enable_checkpointing=False, - model_name="mixtral-8x7b", - dtype="bfloat16", - megablox=True, - sparse_matmul=True, - per_device_batch_size=1, - ici_tensor_transpose_parallelism=4, - max_target_length=128, - ) - - cfg2 = pyconfig.initialize( - [None, get_test_config_path()], - run_name="moe_block_megablox_tp_test", - enable_checkpointing=False, - model_name="mixtral-8x7b", - dtype="bfloat16", - megablox=True, - sparse_matmul=True, - per_device_batch_size=1, - ici_tensor_parallelism=4, - max_target_length=128, - ) - - rng = jax.random.PRNGKey(2345) - rng_model, rng_hidden_states = jax.random.split(rng) - device_count = jax.device_count() - hidden_states = jax.random.uniform( - rng_hidden_states, - (int(cfg.per_device_batch_size) * device_count, cfg.max_target_length, cfg.base_emb_dim), - dtype=cfg.dtype, - ) - - devices_array = maxtext_utils.create_device_mesh(cfg) - mesh = Mesh(devices_array, cfg.mesh_axes) - with nn_partitioning.axis_rules(cfg.logical_axis_rules): - variables, _ = self.get_expected_output(rng_model, hidden_states, cfg, mesh) - tp_transpose_output, _, _ = self.get_moe_output(variables, hidden_states, cfg, mesh) - tp_output, _, _ = self.get_moe_output(variables, hidden_states, cfg2, mesh) - self.assertTrue(jax.numpy.allclose(tp_output, tp_transpose_output, rtol=1e-05, atol=1e-05, equal_nan=False)) - @pytest.mark.tpu_only def test_megablox_context_parallelism(self): cfg = pyconfig.initialize( diff --git a/tests/unit/newly_added_detection_test.py b/tests/unit/newly_added_detection_test.py new file mode 100644 index 0000000000..6c718a4d43 --- /dev/null +++ b/tests/unit/newly_added_detection_test.py @@ -0,0 +1,194 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Unit tests for the line-range test-change detector in ``newly_added_detection``.""" + +import textwrap + +import pytest + +from tests.utils.newly_added_detection import _build_diff_commands +from tests.utils.newly_added_detection import _is_test_file +from tests.utils.newly_added_detection import parse_changed_line_map +from tests.utils.newly_added_detection import touched_test_names + +pytestmark = pytest.mark.cpu_only + + +# --- _is_test_file ----------------------------------------------------------- + + +def test_is_test_file_accepts_test_suffixes_under_tests(): + assert _is_test_file("tests/unit/a_test.py") + assert _is_test_file("tests/integration/b_tests.py") + + +def test_is_test_file_rejects_non_tests_and_non_suffix(): + assert not _is_test_file("tests/utils/test_helpers.py") # helper, not *_test.py + assert not _is_test_file("src/maxtext/foo_test.py") # outside tests/ + assert not _is_test_file("tests/unit/helper.py") # not a test suffix + + +# --- _build_diff_commands ---------------------------------------------------- + + +def test_build_diff_commands_uses_only_merge_base_threedot_ranges(): + # Only three-dot (merge-base) ranges: remote first (CI / local-with-origin), + # then local (developer without an origin remote). Two-dot tip-vs-tip ranges are + # excluded because they over-report every commit the base gained past the fork + # point, dragging unrelated scheduled_only tests into pre-submit. + cmds = _build_diff_commands("main") + assert cmds == [ + ["git", "diff", "--unified=0", "origin/main...HEAD"], + ["git", "diff", "--unified=0", "main...HEAD"], + ] + # Guard against a two-dot form sneaking back in: every range arg is three-dot. + range_args = [cmd[-1] for cmd in cmds] + assert all("..." in arg for arg in range_args) + + +def test_build_diff_commands_honours_non_main_base(): + cmds = _build_diff_commands("release/v2") + assert cmds == [ + ["git", "diff", "--unified=0", "origin/release/v2...HEAD"], + ["git", "diff", "--unified=0", "release/v2...HEAD"], + ] + + +# --- parse_changed_line_map -------------------------------------------------- + + +def test_line_map_added_region_from_header_range(): + diff = ( + "diff --git a/tests/unit/a_test.py b/tests/unit/a_test.py\n" + "--- a/tests/unit/a_test.py\n" + "+++ b/tests/unit/a_test.py\n" + "@@ -0,0 +1,3 @@\n" + "+line1\n+line2\n+line3\n" + ) + assert parse_changed_line_map(diff) == {"tests/unit/a_test.py": {1, 2, 3}} + + +def test_line_map_single_line_default_length(): + diff = "+++ b/tests/unit/a_test.py\n" "@@ -10 +12 @@\n" "+changed\n" + assert parse_changed_line_map(diff) == {"tests/unit/a_test.py": {12}} + + +def test_line_map_pure_deletion_marks_boundary(): + # `+4,0` = pure deletion; the join sits between new-file lines 4 and 5. + diff = "+++ b/tests/unit/a_test.py\n" "@@ -5,2 +4,0 @@\n" "-gone1\n-gone2\n" + assert parse_changed_line_map(diff) == {"tests/unit/a_test.py": {4, 5}} + + +def test_line_map_accumulates_multiple_hunks_and_files(): + diff = ( + "+++ b/tests/unit/a_test.py\n" + "@@ -0,0 +1,1 @@\n" + "+x\n" + "@@ -8,0 +10,2 @@\n" + "+y\n+z\n" + "+++ b/tests/unit/b_test.py\n" + "@@ -0,0 +3,1 @@\n" + "+w\n" + ) + assert parse_changed_line_map(diff) == { + "tests/unit/a_test.py": {1, 10, 11}, + "tests/unit/b_test.py": {3}, + } + + +def test_line_map_ignores_deleted_file(): + diff = "+++ /dev/null\n" "@@ -1,2 +0,0 @@\n" "-a\n-b\n" + assert not parse_changed_line_map(diff) + + +# --- tests_touching_lines ---------------------------------------------------- + +_SOURCE = textwrap.dedent( + """\ + import pytest + + + class TestAlpha: + + @pytest.mark.scheduled_only + def test_existing(self): + x = 1 + assert x == 1 + + def test_untouched(self): + assert True + + + def test_top_level_untouched(): + assert True + + + async def test_async_new(): + assert True + """ +) +# Line numbers in _SOURCE: +# 6 @pytest.mark.scheduled_only +# 7 def test_existing (span 6-9, decorator included) +# 11 def test_untouched (span 11-12) +# 15 def test_top_level_untouched (span 15-16) +# 19 async def test_async_new (span 19-20) + + +def test_touching_body_line_flags_that_test(): + assert touched_test_names(_SOURCE, {9}) == {"test_existing"} + + +def test_touching_decorator_line_flags_that_test(): + assert touched_test_names(_SOURCE, {6}) == {"test_existing"} + + +def test_touching_async_test_is_detected(): + assert touched_test_names(_SOURCE, {20}) == {"test_async_new"} + + +def test_untouched_tests_are_not_flagged(): + # Editing test_existing must not drag in the neighbouring untouched tests. + assert touched_test_names(_SOURCE, {8, 9}) == {"test_existing"} + + +def test_lines_outside_any_test_flag_nothing(): + assert touched_test_names(_SOURCE, {1, 2}) == set() # imports / blank lines + + +def test_empty_touched_set_returns_empty(): + assert touched_test_names(_SOURCE, set()) == set() + + +def test_unparseable_source_returns_empty_gracefully(): + # pytest collection surfaces the syntax error itself; the parser must not raise. + assert touched_test_names("def broken(:\n", {1}) == set() + + +def test_insertion_after_untouched_test_does_not_flag_it(): + # The regression that git's hunk-header heuristic caused: a new test added + # right after an untouched test must flag ONLY the new test. + source = textwrap.dedent( + """\ + def test_old(): + assert True + + + def test_new(): + assert True + """ + ) + # test_old spans lines 1-2; test_new spans lines 5-6. The added lines are 5-6. + assert touched_test_names(source, {5, 6}) == {"test_new"} diff --git a/tests/unit/pallas_mosaic_tpu_v2_kernel_test.py b/tests/unit/pallas_mosaic_tpu_v2_kernel_test.py index 9c605fdbd6..7565838354 100644 --- a/tests/unit/pallas_mosaic_tpu_v2_kernel_test.py +++ b/tests/unit/pallas_mosaic_tpu_v2_kernel_test.py @@ -162,6 +162,7 @@ def reference_tgmm( # group_offset is obtained from # jnp.arange(0, num_experts, num_experts_per_shard) group_offset=None, + partial_sum=None, ): # [num_groups, k, n] """Computes reference transposed grouped matrix multiplication.""" # Compute lhs[:, sizes[i-1]:sizes[i]] @ rhs[sizes[i-1]:sizes[i], :] @@ -179,7 +180,10 @@ def reference_tgmm( group = global_group - group_offset[0] end = start + group_size if 0 <= group < num_actual_groups: - out.append(lhs[:, start:end] @ rhs[start:end, :]) + res = lhs[:, start:end].astype(jnp.float32) @ rhs[start:end, :].astype(jnp.float32) + if partial_sum is not None: + res = res + partial_sum[group].astype(jnp.float32) + out.append(res.astype(lhs.dtype)) start = end return jnp.stack(out) @@ -260,12 +264,13 @@ def test_gmm_basic(self, batch_size, in_size, out_size, num_groups, has_bias, ha in_size=[512, 1024], out_size=[512, 1024], num_groups=[5, 16, 32], + has_partial_sum=[True, False], group_offset=[0, 2, 3], ) - def test_tgmm_basic(self, batch_size, in_size, out_size, num_groups, group_offset): + def test_tgmm_basic(self, batch_size, in_size, out_size, num_groups, has_partial_sum, group_offset): num_local_groups = num_groups - group_offset key = jax.random.key(0) - key1, key2 = jax.random.split(key, 2) + key1, key2, key3 = jax.random.split(key, 3) lhs = jax.random.normal(key1, (batch_size, in_size), dtype=jnp.bfloat16) # [m, k] grad = jax.random.normal(key2, (batch_size, out_size), dtype=jnp.bfloat16) # [m, n] group_sizes = get_group_sizes(batch_size, num_groups) @@ -273,13 +278,18 @@ def test_tgmm_basic(self, batch_size, in_size, out_size, num_groups, group_offse # group_sizes=Array([14, 14, ..., 7]). group_offset = jnp.array(group_offset, dtype=jnp.int32) + ps = None + if has_partial_sum: + ps = jax.random.normal(key3, (num_local_groups, in_size, out_size), dtype=jnp.bfloat16) + lhs_t = lhs.swapaxes(0, 1) # [k, m] - expected = reference_tgmm(lhs_t, grad, group_sizes, num_local_groups, group_offset=group_offset) + expected = reference_tgmm(lhs_t, grad, group_sizes, num_local_groups, group_offset=group_offset, partial_sum=ps) actual = tgmm_backend.tgmm_v2( lhs, grad, group_sizes, num_local_groups, + partial_sum=ps, group_offset=group_offset, preferred_element_type=jnp.bfloat16, ) diff --git a/tests/unit/qwen3_omni_layers_test.py b/tests/unit/qwen3_omni_layers_test.py index 4b3227ed7c..b391dce5f2 100644 --- a/tests/unit/qwen3_omni_layers_test.py +++ b/tests/unit/qwen3_omni_layers_test.py @@ -25,6 +25,7 @@ import jax import jax.numpy as jnp from jax.sharding import Mesh +import pytest from maxtext.configs import pyconfig from maxtext.common import common_types from maxtext.utils.globals import MAXTEXT_REPO_ROOT @@ -37,7 +38,10 @@ ) from maxtext.layers.decoders import deepstack_process from maxtext.layers.encoders import AudioEncoder -from maxtext.multimodal.processor_qwen3_omni import maybe_pad_video_values_to_max_grid +from maxtext.multimodal.processor_qwen3_omni import ( + maybe_pad_video_values_to_max_grid, + preprocess_video, +) from maxtext.models.qwen3 import ( Qwen3OmniAudioEncoder, Qwen3OmniAudioEncoderLayer, @@ -59,7 +63,6 @@ copy_patch_embed_weights, copy_patch_merger_weights, copy_vision_encoder_weights, - create_block_diagonal_attention_mask, create_random_jax_torch, ) import numpy as np @@ -435,6 +438,48 @@ def test_patch_embed_padded_video_valid_outputs_match_unpadded(self): atol=5e-3, ) + def test_scale_to_fit_video_before_padding(self): + """Test that preprocess_video respects video_max_grid_h/w for wide/tall videos.""" + cfg = pyconfig.initialize( + ["", base_config_path], + model_name="qwen3-omni-30b-a3b", + video_max_grid_t=2, + video_max_grid_h=32, + video_max_grid_w=32, + patch_size_for_vit=16, + num_channels_for_vit=3, + temporal_patch_size_for_vit=2, + ) + + # Wide video (2 frames, 3 channels, 220×896): grid_w >> max_grid_w. + dummy_video_wide = np.ones((2, 3, 220, 896), dtype=np.float32) + dummy_video_wide_ratio = dummy_video_wide.shape[3] / dummy_video_wide.shape[2] + video_processed, video_grid_thw = preprocess_video(dummy_video_wide, cfg) + self.assertLessEqual(video_grid_thw[0, 1], cfg.video_max_grid_h) + self.assertLessEqual(video_grid_thw[0, 2], cfg.video_max_grid_w) + # Check the aspect ratio is mostly preserved after smart_resize scaling to fit max grid. + self.assertAlmostEqual(video_grid_thw[0, 2] / video_grid_thw[0, 1], dummy_video_wide_ratio, delta=0.5) + + # Verify maybe_pad_video_values_to_max_grid succeeds without ValueError. + video_values = np.reshape( + video_processed, + ( + 1, + cfg.num_channels_for_vit, + cfg.temporal_patch_size_for_vit * video_grid_thw[0, 0], + cfg.patch_size_for_vit * video_grid_thw[0, 1], + cfg.patch_size_for_vit * video_grid_thw[0, 2], + ), + ) + padded_values, padded_grid, _ = maybe_pad_video_values_to_max_grid( + video_values, + video_grid_thw, + cfg, + ) + self.assertEqual(padded_values.shape[3], cfg.video_max_grid_h * cfg.patch_size_for_vit) + self.assertEqual(padded_values.shape[4], cfg.video_max_grid_w * cfg.patch_size_for_vit) + np.testing.assert_array_equal(padded_grid, video_grid_thw) # Grid should not change + def test_patch_embed_is_jittable(self): """Test that patch embed is JIT-compilable.""" @@ -732,6 +777,7 @@ def test_hidden_states_unchanged_without_visual_tokens(self): np.testing.assert_allclose(np.array(result), hidden_np, rtol=1e-6, atol=1e-6) +@pytest.mark.skip(reason="Requires decord, which may not be installed in remote CI runners.") class TestQwen3OmniPreprocessing(unittest.TestCase): """Test MaxText Qwen3 Omni preprocessor against HuggingFace reference.""" @@ -958,17 +1004,12 @@ def _test_encoder_layer_with_batch_size(self, batch_size): jax_input, torch_input_3d = create_random_jax_torch(batch_size, seq_len, hidden_size) - # PyTorch forward pass - expects 2D input (total_seq_len, hidden_dim) with cu_seqlens - torch_input_2d = torch_input_3d.reshape(-1, hidden_size) - - # Create cu_seqlens for PyTorch (cumulative sequence lengths for each batch) - # For batch_size=2, seq_len=12: [0, 12, 24] indicates two sequences of length 12 each - cu_seqlens = torch.tensor([i * seq_len for i in range(batch_size + 1)], dtype=torch.int32) - - attention_mask = create_block_diagonal_attention_mask(cu_seqlens, torch_input_2d.dtype) - - torch_output_1d = torch_layer(torch_input_2d, cu_seqlens=cu_seqlens, attention_mask=attention_mask)[0] - torch_output = torch_output_1d.reshape(batch_size, seq_len, hidden_size) + # PyTorch audio layers take 2D packed input. Run each batch item separately + # to match MaxText's batched attention without relying on HF's old mask API. + cu_seqlens = torch.tensor([0, seq_len], dtype=torch.int32) + torch_output = torch.stack( + [torch_layer(torch_input_3d[i], cu_seqlens=cu_seqlens)[0] for i in range(batch_size)], dim=0 + ) jax_output = maxtext_layer(jax_input, deterministic=True) @@ -1099,18 +1140,15 @@ def test_audio_encoder_matches_torch(self): torch_after_pos = torch_conv_out + torch_pos_emb # Run through encoder layers + layernorm (but not projector) - # Process all chunks together + # Process chunks separately, matching MaxText's (batch * chunks, seq, hidden) attention shape. seq_len_per_chunk = torch_after_pos.shape[1] - cu_seqlens = torch.tensor([i * seq_len_per_chunk for i in range(num_chunks + 1)], dtype=torch.int32) - attention_mask = create_block_diagonal_attention_mask(cu_seqlens, torch_after_pos.dtype) - - # Flatten: (num_chunks, seq_len_per_chunk, hidden) -> (num_chunks*seq_len_per_chunk, hidden) - hidden_state = torch_after_pos.reshape(-1, torch_after_pos.shape[-1]) + cu_seqlens = torch.tensor([0, seq_len_per_chunk], dtype=torch.int32) + hidden_state = torch_after_pos for layer in torch_model.layers: - hidden_state = layer(hidden_state, cu_seqlens=cu_seqlens, attention_mask=attention_mask)[0] + hidden_state = torch.stack([layer(chunk, cu_seqlens=cu_seqlens)[0] for chunk in hidden_state], dim=0) hidden_state = torch_model.ln_post(hidden_state) - # Reshape back: (num_chunks*seq_len_per_chunk, hidden) -> (batch=1, num_chunks*seq_len_per_chunk, hidden) + # Reshape back: (num_chunks, seq_len_per_chunk, hidden) -> (batch=1, num_chunks*seq_len_per_chunk, hidden) torch_output = hidden_state.reshape(1, num_chunks * seq_len_per_chunk, -1) # MaxText forward diff --git a/tests/unit/tfds_data_processing_test.py b/tests/unit/tfds_data_processing_test.py index fe1bb75f24..c9700d701f 100644 --- a/tests/unit/tfds_data_processing_test.py +++ b/tests/unit/tfds_data_processing_test.py @@ -33,6 +33,7 @@ from tests.utils.test_helpers import get_test_config_path, get_test_base_output_directory +@pytest.mark.cpu_only class TfdsDataProcessingTest(unittest.TestCase): def setUp(self): diff --git a/tests/unit/train_state_nnx_checkpoint_test.py b/tests/unit/train_state_nnx_checkpoint_test.py index fdb94046c7..41e7d19769 100644 --- a/tests/unit/train_state_nnx_checkpoint_test.py +++ b/tests/unit/train_state_nnx_checkpoint_test.py @@ -351,6 +351,22 @@ class TestMaybeSaveCheckpointStepAlignment(unittest.TestCase): def setUp(self): self.tx = optax.adam(1e-3) + def _config(self, **overrides): + """Builds a minimal checkpoint config for maybe_save_checkpoint tests.""" + values = { + "pure_nnx": True, + "checkpoint_period": 10, + "async_checkpointing": False, + "enable_diloco": False, + "enable_continuous_checkpointing": False, + "enable_emergency_checkpoint": False, + "enable_multi_tier_checkpointing": False, + "local_checkpoint_period": 0, + "enable_autocheckpoint": False, + } + values.update(overrides) + return SimpleNamespace(**values) + def _build_nnx_state(self, num_steps): """Build an nnx.State flattened from TrainStateNNX after num_steps gradient applications.""" model = MockModel(rngs=nnx.Rngs(0)) @@ -380,12 +396,7 @@ def _build_linen_state(self, num_steps): def _invoke_maybe_save(self, state, pure_nnx): """Call maybe_save_checkpoint with save_checkpoint patched, return {step, state} captured.""" # checkpoint_period=1 keeps force_ckpt_save False regardless of actual_step. - config = SimpleNamespace( - pure_nnx=pure_nnx, - checkpoint_period=1, - async_checkpointing=False, - enable_diloco=False, - ) + config = self._config(pure_nnx=pure_nnx, checkpoint_period=1) mgr = mock.MagicMock() mgr.reached_preemption.return_value = False @@ -451,12 +462,7 @@ def test_maybe_save_checkpoint_skips_if_already_saved(self): state = self._build_nnx_state(self.N_STEPS) actual_step = self.N_STEPS - 1 - config = SimpleNamespace( - pure_nnx=True, - checkpoint_period=1, - async_checkpointing=False, - enable_diloco=False, - ) + config = self._config(checkpoint_period=1) mgr = mock.MagicMock() mgr.reached_preemption.return_value = False # Mock latest_step to return the same actual_step @@ -475,12 +481,7 @@ def test_maybe_save_checkpoint_saves_if_not_already_saved(self): state = self._build_nnx_state(self.N_STEPS) actual_step = self.N_STEPS - 1 - config = SimpleNamespace( - pure_nnx=True, - checkpoint_period=1, - async_checkpointing=False, - enable_diloco=False, - ) + config = self._config(checkpoint_period=1) mgr = mock.MagicMock() mgr.reached_preemption.return_value = False # Mock latest_step to return a different step (or None) @@ -495,6 +496,101 @@ def test_maybe_save_checkpoint_saves_if_not_already_saved(self): # Assert that save_checkpoint WAS called! save_checkpoint_mock.assert_called_once() + def test_maybe_save_checkpoint_skips_non_checkpoint_step_before_state_work( + self, + ): + """Non-checkpoint steps should not query latest_step or build save args.""" + state = mock.Mock() + config = self._config() + mgr = mock.MagicMock() + mgr.reached_preemption.return_value = False + + with mock.patch.object(checkpointing, "save_checkpoint") as save_checkpoint_mock: + checkpointing.maybe_save_checkpoint(mgr, state, config, data_iterator=None, step=3) + + mgr.latest_step.assert_not_called() + mgr.reached_preemption.assert_called_once_with(3) + mgr.wait_until_finished.assert_not_called() + state.to_pure_dict.assert_not_called() + save_checkpoint_mock.assert_not_called() + + def test_maybe_save_checkpoint_handles_preemption_on_non_checkpoint_step( + self, + ): + """Non-checkpoint steps must still honor preemption handling.""" + state = mock.Mock() + config = self._config() + mgr = mock.MagicMock() + mgr.reached_preemption.return_value = True + + with mock.patch.object(checkpointing, "save_checkpoint") as save_checkpoint_mock: + with self.assertRaises(checkpointing.exceptions.StopTraining): + checkpointing.maybe_save_checkpoint(mgr, state, config, data_iterator=None, step=3) + + mgr.latest_step.assert_not_called() + mgr.reached_preemption.assert_called_once_with(3) + mgr.wait_until_finished.assert_called_once_with() + state.to_pure_dict.assert_not_called() + save_checkpoint_mock.assert_not_called() + + def test_maybe_save_checkpoint_allows_local_checkpoint_period(self): + """Emergency and multi-tier local checkpoint periods dispatch save work.""" + for checkpoint_flag in ( + "enable_emergency_checkpoint", + "enable_multi_tier_checkpointing", + ): + with self.subTest(checkpoint_flag=checkpoint_flag): + state = mock.Mock() + state.to_pure_dict.return_value = { + "model": {}, + "optimizer": {"step": 5}, + } + config = self._config( + checkpoint_period=100, + local_checkpoint_period=5, + **{checkpoint_flag: True}, + ) + mgr = mock.MagicMock() + mgr.latest_step.return_value = None + mgr.reached_preemption.return_value = False + save_checkpoint_mock = mock.MagicMock(return_value=False) + + with mock.patch.object(checkpointing, "save_checkpoint", save_checkpoint_mock): + checkpointing.maybe_save_checkpoint(mgr, state, config, data_iterator=None, step=5) + + mgr.latest_step.assert_called_once_with() + mgr.reached_preemption.assert_called_once_with(5) + mgr.wait_until_finished.assert_not_called() + state.to_pure_dict.assert_called_once_with() + save_checkpoint_mock.assert_called_once() + + def test_maybe_save_checkpoint_allows_mtc_period_with_continuous_policy( + self, + ): + """Continuous checkpointing should not suppress MTC local saves.""" + state = mock.Mock() + state.to_pure_dict.return_value = {"model": {}, "optimizer": {"step": 5}} + config = self._config( + checkpoint_period=100, + enable_continuous_checkpointing=True, + enable_multi_tier_checkpointing=True, + local_checkpoint_period=5, + ) + mgr = mock.MagicMock() + mgr.should_save.return_value = False + mgr.latest_step.return_value = None + mgr.reached_preemption.return_value = False + save_checkpoint_mock = mock.MagicMock(return_value=False) + + with mock.patch.object(checkpointing, "save_checkpoint", save_checkpoint_mock): + checkpointing.maybe_save_checkpoint(mgr, state, config, data_iterator=None, step=5) + + mgr.should_save.assert_called_once_with(5) + mgr.latest_step.assert_called_once_with() + mgr.reached_preemption.assert_called_once_with(5) + state.to_pure_dict.assert_called_once_with() + save_checkpoint_mock.assert_called_once() + class TestLinenCheckpointFormatConverters(unittest.TestCase): """to_linen_checkpoint_dict / from_linen_checkpoint_dict (NNX <-> Linen on-disk layout).""" diff --git a/tests/utils/attention_test_util.py b/tests/utils/attention_test_util.py index 29777d9549..12ff18591a 100644 --- a/tests/utils/attention_test_util.py +++ b/tests/utils/attention_test_util.py @@ -215,8 +215,8 @@ def forward_with_context_expert_parallelism( nn_partitioning.get_axis_rules(), ) pos_spec = nn_partitioning.logical_to_mesh_axes((None, length_axis), nn_partitioning.get_axis_rules()) - lnx_sharding = NamedSharding(mesh_cp, lnx_spec) - pos_sharding = NamedSharding(mesh_cp, pos_spec) + lnx_sharding = NamedSharding(mesh_cp, lnx_spec) # pyrefly: ignore[bad-argument-type] + pos_sharding = NamedSharding(mesh_cp, pos_spec) # pyrefly: ignore[bad-argument-type] lnx = jax.device_put(lnx, lnx_sharding) decoder_segment_ids = jax.device_put(decoder_segment_ids, pos_sharding) diff --git a/tests/utils/forward_pass_logit_checker.py b/tests/utils/forward_pass_logit_checker.py index 92adb4e921..1be6934357 100644 --- a/tests/utils/forward_pass_logit_checker.py +++ b/tests/utils/forward_pass_logit_checker.py @@ -34,7 +34,39 @@ # For example: # tests/assets/logits_generation/golden_llama2-7b_export.ipynb -"""Check if the logits generated by a model's src/maxtext/HF implementation matches golden logits for the same inputs""" +"""Check if the logits generated by a model's src/maxtext/HF implementation matches golden logits for the same inputs. + +Usage: + +1. Comparing MaxText against a pre-generated HF golden logits file +(supported for both text-only and multimodal models; recommended for large models and repeated tests): +For multimodal models, you must pre-generate the HuggingFace golden logits file first using generate_hf_golden_logits.py: + +python3 -m tests.assets.logits_generation.generate_hf_golden_logits \ + --model-id=Qwen/Qwen3-VL-2B-Instruct --hf-model-path=path/to/qwen3-vl-2b \ + --prompts="Describe this image" --image-paths=tests/assets/test_image.jpg \ + --output-path=golden_qwen3-vl-2b_vision.jsonl \ + --apply-chat-template --output-format=json + +Then run the following command to check if the MaxText checkpoint matches the golden logits: + +python3 -m tests.utils.forward_pass_logit_checker src/maxtext/configs/base.yml \ + tokenizer_path=path/to/qwen3-vl-2b \ + load_parameters_path=/qwen3-vl-2b/maxtext_ckpt/0/items model_name=qwen3-vl-2b \ + use_multimodal=true scan_layers=false dtype=float32 per_device_batch_size=1 \ + prompt="Describe this image" image_path=tests/assets/test_image.jpg \ + --max_kl_div=0.1 --golden_logits_path=golden_qwen3-vl-2b_vision.jsonl + +Note: Pre-generating golden logits is also supported and recommended for large text-only models and repeated tests. + +2. Running the HuggingFace model on-the-fly to compare against MaxText (supported for text-only models): + +python3 -m tests.utils.forward_pass_logit_checker src/maxtext/configs/base.yml \ + tokenizer_path=path/to/qwen3-0.6b \ + load_parameters_path=/qwen3-0.6b/maxtext_ckpt/0/items model_name=qwen3-0.6b \ + use_multimodal=false per_device_batch_size=1 dtype=float32 \ + --max_kl_div=0.1 --run_hf_model=true --hf_model_path=path/to/qwen3-0.6b +""" import argparse import functools @@ -169,7 +201,7 @@ def check_kl_divergence(model_logits, golden_logits, atol=0.02): log_target=False, ) - max_logging.log(f"\nAverage KL divergence per token (D_KL(P_golden || Q_model)): {kl_div_value.item():.4e}") + max_logging.log(f"\nAverage KL divergence per token (D_KL(P_golden || Q_model)): {kl_div_value.item():.2e}") # To find the max KL divergence for any single token in the set # use reduction='none'. @@ -180,13 +212,13 @@ def check_kl_divergence(model_logits, golden_logits, atol=0.02): ) # Sum over the vocab dim to get a single KL value per token # Log per-token KL divergences - formatted_list = [f"{x:.4e}" for x in kl_divs_per_token.tolist()] + formatted_list = [f"{x:.2e}" for x in kl_divs_per_token.tolist()] max_logging.log(f"Per-token KL Divergences: \n{formatted_list}") max_kl_div = kl_divs_per_token.max() - max_logging.log(f"\nMax KL divergence for a single token in the set: {max_kl_div.item():.4e}") + max_logging.log(f"\nMax KL divergence for a single token in the set: {max_kl_div.item():.2e}") - assert max_kl_div < atol, f"KL divergence values {max_kl_div.item():.4e} exceed the threshold {atol}" + assert max_kl_div < atol, f"KL divergence values {max_kl_div.item():.2e} exceed the threshold {atol}" def get_data(golden_data_point, config): @@ -211,6 +243,16 @@ def get_data(golden_data_point, config): pixel_values = np.transpose(pixel_values, (1, 2, 0)) elif model_prefix in ["llama4"]: pixel_values = pixel_values[None, :] + elif model_prefix in ["qwen3"]: + grid_thw = np.asarray(golden_data_point["image_grid_thw"]) + if grid_thw.ndim == 2: + grid_t, grid_h, grid_w = grid_thw[0] + else: + grid_t, grid_h, grid_w = grid_thw + tps = config.temporal_patch_size_for_vit + p = config.patch_size_for_vit + c = config.num_channels_for_vit + pixel_values = np.reshape(pixel_values, (c, int(grid_t * tps), int(grid_h * p), int(grid_w * p))) pixel_values = np.stack([pixel_values for _ in range(config.global_batch_size_to_train_on)]) else: pixel_values = None @@ -234,13 +276,35 @@ def get_data(golden_data_point, config): prompt = golden_data_point["formatted_prompt"] else: prompt = golden_data_point["prompt"] - max_logging.log(f' prompt="{prompt}" raw ids={original_ids}, logits.shape = {logits.shape}') + max_logging.log(f' prompt="{prompt}" raw ids (first 20)={original_ids[:20]}, logits.shape = {logits.shape}') decoder_segment_ids = np.zeros(s, dtype=np.int32) decoder_segment_ids[:, :seq_len] = DECODING_ACTIVE_SEQUENCE_INDICATOR - decoder_positions = np.stack( - [np.arange(config.max_target_length, dtype=np.int32) for _ in range(config.global_batch_size_to_train_on)] - ) + + # For Qwen3-VL model, compute RoPE embeddings to generate 3D position IDs + if model_prefix in ["qwen3"] and config.use_mrope: + from maxtext.multimodal import processor_qwen3_omni # pylint: disable=import-outside-toplevel + + image_grid_thw = np.atleast_2d(golden_data_point["image_grid_thw"]) if "image_grid_thw" in golden_data_point else None + video_grid_thw = np.atleast_2d(golden_data_point["video_grid_thw"]) if "video_grid_thw" in golden_data_point else None + + attention_mask = np.zeros(s, dtype=np.int32) + attention_mask[:, :seq_len] = 1 + + position_ids, _ = processor_qwen3_omni.get_rope_index( + input_ids=ids, + image_grid_thw=image_grid_thw, + video_grid_thw=video_grid_thw, + attention_mask=attention_mask, + spatial_merge_size=config.spatial_merge_size_for_vit, + position_id_per_seconds=config.position_id_per_seconds, + config=config, + ) + decoder_positions = position_ids.astype(np.int32) + else: # For non-Qwen3-VL models, keep 1D position IDs + decoder_positions = np.stack( + [np.arange(config.max_target_length, dtype=np.int32) for _ in range(config.global_batch_size_to_train_on)] + ) return ids, decoder_segment_ids, decoder_positions, logits, seq_len, pixel_values @@ -251,6 +315,30 @@ def main(config, test_args): # pylint: disable=W0621 devices_array = maxtext_utils.create_device_mesh(config) mesh = jax.sharding.Mesh(devices_array, config.mesh_axes) + # Load tokenizer so it is available for: + # 1. Pre-loaded golden logits comparison (multimodal input) + # 2. On-the-fly HuggingFace model comparison (text only input) + hf_token = config.hf_access_token + try: + if test_args.hf_model_path: + max_logging.log(f"Loading tokenizer from {test_args.hf_model_path}.") + tokenizer = AutoTokenizer.from_pretrained( + test_args.hf_model_path, token=hf_token, trust_remote_code=test_args.trust_remote_code + ) + else: + max_logging.log(f"Loading tokenizer from {config.tokenizer_path}.") + tokenizer = AutoTokenizer.from_pretrained( + config.tokenizer_path, token=hf_token, trust_remote_code=test_args.trust_remote_code + ) + except Exception as e: # pylint: disable=broad-except + max_logging.log(f"Tokenizer loading error: {e}.\nLoading tokenizer from {config.tokenizer_path}.") + tokenizer = AutoTokenizer.from_pretrained( + config.tokenizer_path, token=hf_token, trust_remote_code=test_args.trust_remote_code + ) + + if config.model_name.startswith(("llama3.1", "mixtral")): + tokenizer.pad_token = tokenizer.eos_token + if not test_args.run_hf_model: """Comparing maxtext/huggingface model with pre-loaded golden logitis""" max_logging.log("Initializing MaxText model") @@ -261,8 +349,7 @@ def main(config, test_args): # pylint: disable=W0621 if config.lora.enable_lora: model = lora_utils.apply_lora_to_model(model, mesh, config) if config.lora.lora_restore_path: - mock_trainer = type("MockTrainer", (), {"model": model, "train_steps": 0}) - lora_utils.restore_lora_from_path(mock_trainer, config) + lora_utils.restore_lora_from_path(model, config) state = None else: model = models.transformer_as_linen(config, mesh=mesh, quant=quant, model_mode=MODEL_MODE_TRAIN) @@ -351,13 +438,29 @@ def main(config, test_args): # pylint: disable=W0621 max_rel_diff_val = rel_diff[max_rel_diff_idx] msg = ( "\n[numerical difference]\n" - f"Max absolute difference: {max_abs_diff_val:.4e} at index {max_abs_diff_idx}\n" - f" (Train: {train_logits_slice[max_abs_diff_idx]:.4e}, Golden: {golden_logits_slice[max_abs_diff_idx]:.4e})\n" - f"Max relative difference: {max_rel_diff_val:.4e} at index {max_rel_diff_idx}\n" - f" (Train: {train_logits_slice[max_rel_diff_idx]:.4e}, Golden: {golden_logits_slice[max_rel_diff_idx]:.4e})" + f"Max absolute difference: {max_abs_diff_val:.2e} at index {max_abs_diff_idx}\n" + f" (Train: {train_logits_slice[max_abs_diff_idx]:.2e}, Golden: {golden_logits_slice[max_abs_diff_idx]:.2e})\n" + f"Max relative difference: {max_rel_diff_val:.2e} at index {max_rel_diff_idx}\n" + f" (Train: {train_logits_slice[max_rel_diff_idx]:.2e}, Golden: {golden_logits_slice[max_rel_diff_idx]:.2e})" ) max_logging.log(msg) + # --- Compare logits for top-k token comparison table --- + mt_last_token_logits = ( + convert_jax_weight_to_torch(train_logits_slice[-1:, :]) + if isinstance(train_logits_slice, jax.Array) + else torch.tensor(train_logits_slice[-1:, :]) + ) + hf_last_token_logits = ( + convert_jax_weight_to_torch(golden_logits_slice[-1:, :]) + if isinstance(golden_logits_slice, jax.Array) + else torch.tensor(golden_logits_slice[-1:, :]) + ) + + tokens_maxtext = get_top_k_tokens_scores(mt_last_token_logits, tokenizer, k=10, description="MaxText model") + tokens_hf = get_top_k_tokens_scores(hf_last_token_logits, tokenizer, k=10, description="HF model") + compare_top_tokens(converted_tokens=tokens_maxtext, golden_tokens=tokens_hf) + if test_args.clip_logits_epsilon is not None: model_probabilities = jnp.clip(jax.nn.softmax(train_logits_slice, axis=-1), min=test_args.clip_logits_epsilon) golden_probabilities = jnp.clip(jax.nn.softmax(golden_logits_slice, axis=-1), min=test_args.clip_logits_epsilon) @@ -371,17 +474,41 @@ def main(config, test_args): # pylint: disable=W0621 max_logging.log(f"{model_probabilities[1]=}") kl_div = jax.numpy.sum(jax.scipy.special.kl_div(golden_probabilities, model_probabilities), axis=-1) + + # Mask out vision placeholder tokens for KL calculation + ignore_token_ids = [] + if "qwen3" in config.model_name.lower(): + from maxtext.multimodal.processor_qwen3_omni import QwenTokens # pylint: disable=import-outside-toplevel + + qwen_tokens = QwenTokens(config) + ignore_token_ids = [ + qwen_tokens.vision_start, + qwen_tokens.vision_end, + qwen_tokens.image_pad, + qwen_tokens.video_pad, + ] + + if ignore_token_ids: + slice_ids = ids[0, start_index:token_size] + mask = jnp.ones_like(slice_ids, dtype=jnp.bool_) + for ignore_id in ignore_token_ids: + mask = mask & (slice_ids != ignore_id) + kl_div = jnp.where(mask, kl_div, 0.0) + max_kl_div_val = jax.numpy.max(kl_div) max_kl_div_idx = jax.numpy.argmax(kl_div) + + # print per token kl for debugging, omit for multimodal as the seq length is usually long + per_token_kl = "" if config.use_multimodal else f"KL divergence = {kl_div}, " max_logging.log( f"\n[KL divergence]\n" - f"KL divergence = {kl_div}, max KL divergence = {max_kl_div_val} at index {max_kl_div_idx}, " + f"{per_token_kl} max KL divergence = {max_kl_div_val} at index {max_kl_div_idx}, " f"the corresponding token id is {ids[0, max_kl_div_idx + start_index]}" ) if jax.process_index() == 0 and test_args.output_logits_path: data_to_save = { - "prompt": golden_data[golden_data_index]["prompt"], + "prompt": golden_data_point["prompt"], "tokens": ids[0, :seq_len].tolist(), "logits": full_train_logits[0].tolist(), } @@ -428,11 +555,18 @@ def main(config, test_args): # pylint: disable=W0621 } # Default to bfloat16 if dtype is unrecognized - torch_dtype = dtype_mapping.get(config.dtype.name, torch.bfloat16) + torch_dtype = dtype_mapping.get(config.dtype.name.lower(), torch.bfloat16) max_logging.log(f"Loading HF model with dtype: {torch_dtype} (derived from config.dtype: {config.dtype})") - hf_model = AutoModelForCausalLM.from_pretrained( - test_args.hf_model_path, dtype=torch_dtype, token=hf_token, trust_remote_code=test_args.trust_remote_code + if "qwen3-vl" in config.model_name.lower(): + from transformers import Qwen3VLForConditionalGeneration # pylint: disable=import-outside-toplevel + + model_class = Qwen3VLForConditionalGeneration + else: + model_class = AutoModelForCausalLM + + hf_model = model_class.from_pretrained( + test_args.hf_model_path, torch_dtype=torch_dtype, token=hf_token, trust_remote_code=test_args.trust_remote_code ) hf_lora_path = config.hf_lora_adapter_path if hf_lora_path: @@ -443,25 +577,6 @@ def main(config, test_args): # pylint: disable=W0621 raise ImportError("peft library is required to load HF LoRA adapter. Run `pip install peft`.") from exc hf_model = PeftModel.from_pretrained(hf_model, hf_lora_path) - # Load tokenizer: `test_args.hf_model_path` or fallback to `config.tokenizer_path` - try: - # Try loading from `test_args.hf_model_path` - max_logging.log(f"Loading tokenizer from {test_args.hf_model_path}.") - tokenizer = AutoTokenizer.from_pretrained( - test_args.hf_model_path, token=hf_token, trust_remote_code=test_args.trust_remote_code - ) - except Exception as e: # pylint: disable=broad-except - # Fallback to `config.tokenizer_path`. local hf directory may not contain tokenizer, read from remote tokenizer - max_logging.log(f"Tokenizer loading error: {e}.\nLoading tokenizer from {config.tokenizer_path}.") - tokenizer = AutoTokenizer.from_pretrained( - config.tokenizer_path, token=hf_token, trust_remote_code=test_args.trust_remote_code - ) - - # maxtext model prefix, use eos token as pad token - pad_token_prefixes = ["llama3.1", "mixtral"] - if any(config.model_name.startswith(prefix) for prefix in pad_token_prefixes): - tokenizer.pad_token = tokenizer.eos_token - quant = quantizations.configure_quantization(config) if config.pure_nnx_decoder and config.enable_nnx: maxtext_model = model_creation_utils.from_pretrained(config, mesh=mesh, model_mode=MODEL_MODE_TRAIN) @@ -469,8 +584,7 @@ def main(config, test_args): # pylint: disable=W0621 if config.lora.enable_lora: maxtext_model = lora_utils.apply_lora_to_model(maxtext_model, mesh, config) if config.lora.lora_restore_path: - mock_trainer = type("MockTrainer", (), {"model": maxtext_model, "train_steps": 0}) - lora_utils.restore_lora_from_path(mock_trainer, config) + lora_utils.restore_lora_from_path(maxtext_model, config) maxtext_state = None else: maxtext_model = models.transformer_as_linen(config, mesh, quant=quant, model_mode=MODEL_MODE_TRAIN) diff --git a/tests/utils/newly_added_detection.py b/tests/utils/newly_added_detection.py new file mode 100644 index 0000000000..314dcaaae4 --- /dev/null +++ b/tests/utils/newly_added_detection.py @@ -0,0 +1,206 @@ +# Copyright 2026 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# https://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Detect the tests a PR adds or modifies, for scheduled_only pre-submit verification. + +The pre-submit CI runs ``pytest -m " and (not scheduled_only or newly_added)"``. +This module supplies the ``newly_added`` set: the tests a pull request touched, so that a +newly added or modified ``scheduled_only`` test runs at least once before merge instead of +being silently skipped until the nightly scheduled pipeline. + +Detection maps changed *line numbers* (from ``git diff --unified=0``) onto each test's +line span (from an ``ast`` parse of the new file). A test counts as changed only when a +changed line lands inside its own span. This avoids trusting git's hunk-header function +name, which points at the function *preceding* an insertion and would otherwise flag an +untouched test that merely sits above newly added code. + +Only the Python standard library is used, since this runs on a bare CI runner where +MaxText is not necessarily importable. +""" + +import ast +import os +import re +import subprocess + +# Matches pytest.ini's ``python_files = *_test.py *_tests.py``. +_TEST_SUFFIXES = ("_test.py", "_tests.py") +# Matches pytest.ini's ``testpaths = tests``. +_TESTS_ROOT = "tests/" +# Captures the new-file start line and (optional) length from a unified-diff hunk header: +# @@ -[,] +[,] @@ +_HUNK_RE = re.compile(r"^@@ -\d+(?:,\d+)? \+(\d+)(?:,(\d+))? @@") + + +def _is_test_file(path): + """Return True if ``path`` is a MaxText test file (under tests/ with a test suffix).""" + normalized = path.strip().replace(os.sep, "/") + return normalized.startswith(_TESTS_ROOT) and normalized.endswith(_TEST_SUFFIXES) + + +def parse_changed_line_map(diff_text): + """Map each changed file to the set of new-file line numbers it touched. + + Args: + diff_text: Raw ``git diff --unified=0`` output. + + Returns: + A dict of ``{file_path: set_of_new_line_numbers}``. ``file_path`` is repo-root + relative (the ``b/`` prefix is stripped). For a pure deletion (new length 0) the + two lines bracketing the removal are recorded, so a test that had lines removed is + still detected. Deleted files (``+++ /dev/null``) are omitted. + """ + line_map = {} + current_file = None + for line in diff_text.splitlines(): + if line.startswith("+++"): + path = line[3:].strip() + if path.startswith("b/"): + path = path[2:] + current_file = None if path == "/dev/null" else path + continue + if line.startswith("@@"): + if current_file is None: + continue + match = _HUNK_RE.match(line) + if match is None: + continue + new_start = int(match.group(1)) + new_len = int(match.group(2)) if match.group(2) is not None else 1 + touched = line_map.setdefault(current_file, set()) + if new_len > 0: + touched.update(range(new_start, new_start + new_len)) + else: + touched.update({new_start, new_start + 1}) + return line_map + + +def _iter_test_defs(tree): + """Yield ``(name, start_line, end_line)`` for every test function in an AST. + + Covers module-level test functions and methods declared directly inside a class, + which is what pytest collects. The span starts at the first decorator (if any) so + decorator-only edits are attributed to the test they decorate. + """ + def_types = (ast.FunctionDef, ast.AsyncFunctionDef) + for node in tree.body: + if isinstance(node, def_types) and node.name.startswith("test_"): + yield node.name, _span_start(node), node.end_lineno + elif isinstance(node, ast.ClassDef): + for sub in node.body: + if isinstance(sub, def_types) and sub.name.startswith("test_"): + yield sub.name, _span_start(sub), sub.end_lineno + + +def _span_start(node): + """Return the first source line of a def, including any decorator lines above it.""" + start = node.lineno + if node.decorator_list: + start = min(start, min(dec.lineno for dec in node.decorator_list)) + return start + + +def touched_test_names(source, touched_lines): + """Return the names of test functions in ``source`` whose span includes a changed line. + + Args: + source: The new file's Python source. + touched_lines: Set of changed new-file line numbers for that file. + + Returns: + A set of test function names. Empty if ``touched_lines`` is empty or ``source`` does + not parse (pytest collection surfaces a syntax error on its own, so raising here would + only hide it). + """ + if not touched_lines: + return set() + try: + tree = ast.parse(source) + except SyntaxError: + return set() + found = set() + for name, start, end in _iter_test_defs(tree): + if any(start <= line <= end for line in touched_lines): + found.add(name) + return found + + +def _build_diff_commands(base): + """Return the ordered ``git diff`` argument lists to try, most-precise first. + + Both ranges are three-dot (merge-base) ranges, so only the branch's own commits + are reported. Two-dot tip-vs-tip ranges are deliberately excluded: they over-report + every commit the base gained past the fork point (a stale local ``main`` can inflate + the changed set many-fold), which would drag unrelated ``scheduled_only`` tests into + pre-submit. The remote range covers CI and local checkouts that have an ``origin`` + remote; the local range is the fallback for a developer without ``origin`` and is + never reached in CI, where ``origin/`` is always fetched first. + """ + return [ + ["git", "diff", "--unified=0", f"origin/{base}...HEAD"], + ["git", "diff", "--unified=0", f"{base}...HEAD"], + ] + + +def get_changed_tests(base_ref=None): + """Return ``(file_path, test_name)`` for every test the PR added or modified. + + Args: + base_ref: Base git ref to diff against. Defaults to the ``GITHUB_BASE_REF`` + environment variable, then ``"main"``. + + Returns: + A set of ``(file_path, test_name)`` tuples, or an empty set when not in a git work + tree or when the diff cannot be computed. + """ + base = base_ref or os.environ.get("GITHUB_BASE_REF") or "main" + try: + inside = subprocess.run( + ["git", "rev-parse", "--is-inside-work-tree"], + capture_output=True, + check=False, + ) + if inside.returncode != 0: + return set() + if os.environ.get("GITHUB_ACTIONS") == "true": + subprocess.run( + ["git", "fetch", "origin", f"{base}:refs/remotes/origin/{base}"], + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + check=False, + ) + diff_text = None + for command in _build_diff_commands(base): + try: + diff_text = subprocess.check_output(command, text=True, stderr=subprocess.DEVNULL) + break + except Exception: # pylint: disable=broad-exception-caught + continue + if diff_text is None: + return set() + except Exception: # pylint: disable=broad-exception-caught + return set() + + changed = set() + for path, touched_lines in parse_changed_line_map(diff_text).items(): + if not _is_test_file(path): + continue + try: + with open(path, "r", encoding="utf-8") as handle: + source = handle.read() + except OSError: + continue + for name in touched_test_names(source, touched_lines): + changed.add((path, name)) + return changed diff --git a/tests/utils/reference_hlo_deepseek3.txt b/tests/utils/reference_hlo_deepseek3.txt index eb316f1963..c81d52bf7d 100644 --- a/tests/utils/reference_hlo_deepseek3.txt +++ b/tests/utils/reference_hlo_deepseek3.txt @@ -1,4 +1,4 @@ -HloModule jit_train_step, is_scheduled=true, input_output_alias={ {0}: (0, {}, may-alias), {1}: (1, {}, may-alias), {2}: (2, {}, may-alias), {3}: (3, {}, may-alias), {4}: (4, {}, may-alias), {5}: (5, {}, may-alias), {6}: (6, {}, may-alias), {7}: (7, {}, may-alias), {8}: (8, {}, may-alias), {9}: (9, {}, may-alias), {10}: (10, {}, may-alias), {11}: (11, {}, may-alias), {12}: (12, {}, may-alias), {13}: (13, {}, may-alias), {14}: (14, {}, may-alias), {15}: (15, {}, may-alias), {16}: (16, {}, may-alias), {17}: (17, {}, may-alias), {18}: (18, {}, may-alias), {19}: (19, {}, may-alias), {20}: (20, {}, may-alias), {21}: (21, {}, may-alias), {22}: (22, {}, may-alias), {23}: (23, {}, may-alias), {24}: (24, {}, may-alias), {25}: (25, {}, may-alias), {26}: (26, {}, may-alias), {27}: (27, {}, may-alias), {28}: (28, {}, may-alias), {29}: (29, {}, may-alias), {30}: (30, {}, may-alias), {31}: (31, {}, may-alias), {32}: (32, {}, may-alias), {33}: (33, {}, may-alias), {34}: (34, {}, may-alias), {35}: (35, {}, may-alias), {36}: (36, {}, may-alias), {37}: (37, {}, may-alias), {38}: (38, {}, may-alias), {39}: (39, {}, may-alias), {40}: (40, {}, may-alias), {41}: (41, {}, may-alias), {42}: (42, {}, may-alias), {43}: (43, {}, may-alias), {44}: (44, {}, may-alias), {45}: (45, {}, may-alias), {46}: (46, {}, may-alias), {47}: (47, {}, may-alias), {48}: (48, {}, may-alias), {49}: (49, {}, may-alias), {50}: (50, {}, may-alias), {51}: (51, {}, may-alias), {52}: (52, {}, may-alias), {53}: (53, {}, may-alias), {54}: (54, {}, may-alias), {55}: (55, {}, may-alias), {56}: (56, {}, may-alias), {57}: (57, {}, may-alias), {58}: (58, {}, may-alias), {59}: (59, {}, may-alias), {60}: (60, {}, may-alias), {61}: (61, {}, may-alias), {62}: (62, {}, may-alias), {63}: (63, {}, may-alias), {64}: (64, {}, may-alias), {65}: (65, {}, may-alias), {66}: (66, {}, may-alias), {67}: (67, {}, may-alias), {68}: (68, {}, may-alias), {69}: (69, {}, may-alias), {70}: (70, {}, may-alias), {71}: (71, {}, may-alias), {72}: (72, {}, may-alias), {73}: (73, {}, may-alias), {74}: (74, {}, may-alias), {75}: (75, {}, may-alias), {76}: (76, {}, may-alias), {77}: (77, {}, may-alias), {78}: (78, {}, may-alias), {79}: (79, {}, may-alias), {80}: (80, {}, may-alias), {81}: (81, {}, may-alias), {82}: (82, {}, may-alias), {83}: (83, {}, may-alias), {84}: (84, {}, may-alias), {85}: (85, {}, may-alias), {86}: (86, {}, may-alias), {87}: (87, {}, may-alias), {88}: (88, {}, may-alias), {89}: (89, {}, may-alias), {90}: (90, {}, may-alias), {91}: (91, {}, may-alias), {92}: (92, {}, may-alias), {93}: (93, {}, may-alias), {94}: (94, {}, may-alias), {95}: (95, {}, may-alias), {96}: (96, {}, may-alias), {97}: (97, {}, may-alias), {98}: (98, {}, may-alias) }, entry_computation_layout={(s32[]{:T(128)}, f32[512]{0:T(512)}, f32[512,3,18432]{2,0,1:T(8,128)}, f32[512,3,18432]{2,0,1:T(8,128)}, f32[18432,3,512]{2,0,1:T(8,128)}, /*index=5*/f32[512,3]{0,1:T(4,128)}, f32[512,3]{0,1:T(4,128)}, f32[512,3]{0,1:T(4,128)}, f32[128,3,128,512]{3,2,1,0:T(8,128)}, f32[1536,3]{0,1:T(4,128)}, /*index=10*/f32[512,3,576]{0,2,1:T(8,128)}, f32[512,3,128,256]{3,2,1,0:T(8,128)}, f32[512,3,1536]{2,0,1:T(8,128)}, f32[1536,3,128,192]{2,3,1,0:T(8,128)}, f32[512,129280]{1,0:T(8,128)}, /*index=15*/f32[256,1]{0,1:T(1,128)}, f32[512,1,256]{2,1,0:T(1,128)}, f32[256,1,512,512]{3,2,1,0:T(8,128)}, f32[256,1,512,512]{3,2,1,0:T(8,128)}, f32[256,1,512,512]{3,2,1,0:T(8,128)}, /*index=20*/f32[512,1,512]{2,1,0:T(1,128)}, f32[512,1,512]{2,1,0:T(1,128)}, f32[512,1,512]{2,1,0:T(1,128)}, f32[512,1]{0,1:T(1,128)}, f32[512,1]{0,1:T(1,128)}, /*index=25*/f32[512,1]{0,1:T(1,128)}, f32[128,1,128,512]{3,2,1,0:T(8,128)}, f32[1536,1]{0,1:T(1,128)}, f32[512,1,576]{0,2,1:T(8,128)}, f32[512,1,128,256]{3,2,1,0:T(8,128)}, /*index=30*/f32[512,1,1536]{2,1,0:T(1,128)}, f32[1536,1,128,192]{2,3,1,0:T(8,128)}, f32[129280,512]{1,0:T(8,128)}, s32[]{:T(128)}, f32[512]{0:T(512)}, /*index=35*/f32[512,3,18432]{2,0,1:T(8,128)}, f32[512,3,18432]{2,0,1:T(8,128)}, f32[18432,3,512]{2,0,1:T(8,128)}, f32[512,3]{0,1:T(4,128)}, f32[512,3]{0,1:T(4,128)}, /*index=40*/f32[512,3]{0,1:T(4,128)}, f32[128,3,128,512]{3,2,1,0:T(8,128)}, f32[1536,3]{0,1:T(4,128)}, f32[512,3,576]{0,2,1:T(8,128)}, f32[512,3,128,256]{3,2,1,0:T(8,128)}, /*index=45*/f32[512,3,1536]{2,0,1:T(8,128)}, f32[1536,3,128,192]{2,3,1,0:T(8,128)}, f32[512,129280]{1,0:T(8,128)}, f32[256,1]{0,1:T(1,128)}, f32[512,1,256]{2,1,0:T(1,128)}, /*index=50*/f32[256,1,512,512]{3,2,1,0:T(8,128)}, f32[256,1,512,512]{3,2,1,0:T(8,128)}, f32[256,1,512,512]{3,2,1,0:T(8,128)}, f32[512,1,512]{2,1,0:T(1,128)}, f32[512,1,512]{2,1,0:T(1,128)}, /*index=55*/f32[512,1,512]{2,1,0:T(1,128)}, f32[512,1]{0,1:T(1,128)}, f32[512,1]{0,1:T(1,128)}, f32[512,1]{0,1:T(1,128)}, f32[128,1,128,512]{3,2,1,0:T(8,128)}, /*index=60*/f32[1536,1]{0,1:T(1,128)}, f32[512,1,576]{0,2,1:T(8,128)}, f32[512,1,128,256]{3,2,1,0:T(8,128)}, f32[512,1,1536]{2,1,0:T(1,128)}, f32[1536,1,128,192]{2,3,1,0:T(8,128)}, /*index=65*/f32[129280,512]{1,0:T(8,128)}, f32[512]{0:T(512)}, f32[512,3,18432]{2,0,1:T(8,128)}, f32[512,3,18432]{2,0,1:T(8,128)}, f32[18432,3,512]{2,0,1:T(8,128)}, /*index=70*/f32[512,3]{0,1:T(4,128)}, f32[512,3]{0,1:T(4,128)}, f32[512,3]{0,1:T(4,128)}, f32[128,3,128,512]{3,2,1,0:T(8,128)}, f32[1536,3]{0,1:T(4,128)}, /*index=75*/f32[512,3,576]{0,2,1:T(8,128)}, f32[512,3,128,256]{3,2,1,0:T(8,128)}, f32[512,3,1536]{2,0,1:T(8,128)}, f32[1536,3,128,192]{2,3,1,0:T(8,128)}, f32[512,129280]{1,0:T(8,128)}, /*index=80*/f32[256,1]{0,1:T(1,128)}, f32[512,1,256]{2,1,0:T(1,128)}, f32[256,1,512,512]{3,2,1,0:T(8,128)}, f32[256,1,512,512]{3,2,1,0:T(8,128)}, f32[256,1,512,512]{3,2,1,0:T(8,128)}, /*index=85*/f32[512,1,512]{2,1,0:T(1,128)}, f32[512,1,512]{2,1,0:T(1,128)}, f32[512,1,512]{2,1,0:T(1,128)}, f32[512,1]{0,1:T(1,128)}, f32[512,1]{0,1:T(1,128)}, /*index=90*/f32[512,1]{0,1:T(1,128)}, f32[128,1,128,512]{3,2,1,0:T(8,128)}, f32[1536,1]{0,1:T(1,128)}, f32[512,1,576]{0,2,1:T(8,128)}, f32[512,1,128,256]{3,2,1,0:T(8,128)}, /*index=95*/f32[512,1,1536]{2,1,0:T(1,128)}, f32[1536,1,128,192]{2,3,1,0:T(8,128)}, f32[129280,512]{1,0:T(8,128)}, s32[]{:T(128)}, s32[4,128]{1,0:T(4,128)}, /*index=100*/s32[4,128]{1,0:T(4,128)}, s32[4,128]{1,0:T(4,128)}, s32[4,128]{1,0:T(4,128)}, s32[4,128]{1,0:T(4,128)})->(s32[]{:T(128)}, f32[512]{0:T(512)}, f32[512,3,18432]{2,0,1:T(8,128)}, f32[512,3,18432]{2,0,1:T(8,128)}, f32[18432,3,512]{2,0,1:T(8,128)}, /*index=5*/f32[512,3]{0,1:T(4,128)}, f32[512,3]{0,1:T(4,128)}, f32[512,3]{0,1:T(4,128)}, f32[128,3,128,512]{3,2,1,0:T(8,128)}, f32[1536,3]{0,1:T(4,128)}, /*index=10*/f32[512,3,576]{0,2,1:T(8,128)}, f32[512,3,128,256]{3,2,1,0:T(8,128)}, f32[512,3,1536]{2,0,1:T(8,128)}, f32[1536,3,128,192]{2,3,1,0:T(8,128)}, f32[512,129280]{1,0:T(8,128)}, /*index=15*/f32[256,1]{0,1:T(1,128)}, f32[512,1,256]{2,1,0:T(1,128)}, f32[256,1,512,512]{3,2,1,0:T(8,128)}, f32[256,1,512,512]{3,2,1,0:T(8,128)}, f32[256,1,512,512]{3,2,1,0:T(8,128)}, /*index=20*/f32[512,1,512]{2,1,0:T(1,128)}, f32[512,1,512]{2,1,0:T(1,128)}, f32[512,1,512]{2,1,0:T(1,128)}, f32[512,1]{0,1:T(1,128)}, f32[512,1]{0,1:T(1,128)}, /*index=25*/f32[512,1]{0,1:T(1,128)}, f32[128,1,128,512]{3,2,1,0:T(8,128)}, f32[1536,1]{0,1:T(1,128)}, f32[512,1,576]{0,2,1:T(8,128)}, f32[512,1,128,256]{3,2,1,0:T(8,128)}, /*index=30*/f32[512,1,1536]{2,1,0:T(1,128)}, f32[1536,1,128,192]{2,3,1,0:T(8,128)}, f32[129280,512]{1,0:T(8,128)}, s32[]{:T(128)}, f32[512]{0:T(512)}, /*index=35*/f32[512,3,18432]{2,0,1:T(8,128)}, f32[512,3,18432]{2,0,1:T(8,128)}, f32[18432,3,512]{2,0,1:T(8,128)}, f32[512,3]{0,1:T(4,128)}, f32[512,3]{0,1:T(4,128)}, /*index=40*/f32[512,3]{0,1:T(4,128)}, f32[128,3,128,512]{3,2,1,0:T(8,128)}, f32[1536,3]{0,1:T(4,128)}, f32[512,3,576]{0,2,1:T(8,128)}, f32[512,3,128,256]{3,2,1,0:T(8,128)}, /*index=45*/f32[512,3,1536]{2,0,1:T(8,128)}, f32[1536,3,128,192]{2,3,1,0:T(8,128)}, f32[512,129280]{1,0:T(8,128)}, f32[256,1]{0,1:T(1,128)}, f32[512,1,256]{2,1,0:T(1,128)}, /*index=50*/f32[256,1,512,512]{3,2,1,0:T(8,128)}, f32[256,1,512,512]{3,2,1,0:T(8,128)}, f32[256,1,512,512]{3,2,1,0:T(8,128)}, f32[512,1,512]{2,1,0:T(1,128)}, f32[512,1,512]{2,1,0:T(1,128)}, /*index=55*/f32[512,1,512]{2,1,0:T(1,128)}, f32[512,1]{0,1:T(1,128)}, f32[512,1]{0,1:T(1,128)}, f32[512,1]{0,1:T(1,128)}, f32[128,1,128,512]{3,2,1,0:T(8,128)}, /*index=60*/f32[1536,1]{0,1:T(1,128)}, f32[512,1,576]{0,2,1:T(8,128)}, f32[512,1,128,256]{3,2,1,0:T(8,128)}, f32[512,1,1536]{2,1,0:T(1,128)}, f32[1536,1,128,192]{2,3,1,0:T(8,128)}, /*index=65*/f32[129280,512]{1,0:T(8,128)}, f32[512]{0:T(512)}, f32[512,3,18432]{2,0,1:T(8,128)}, f32[512,3,18432]{2,0,1:T(8,128)}, f32[18432,3,512]{2,0,1:T(8,128)}, /*index=70*/f32[512,3]{0,1:T(4,128)}, f32[512,3]{0,1:T(4,128)}, f32[512,3]{0,1:T(4,128)}, f32[128,3,128,512]{3,2,1,0:T(8,128)}, f32[1536,3]{0,1:T(4,128)}, /*index=75*/f32[512,3,576]{0,2,1:T(8,128)}, f32[512,3,128,256]{3,2,1,0:T(8,128)}, f32[512,3,1536]{2,0,1:T(8,128)}, f32[1536,3,128,192]{2,3,1,0:T(8,128)}, f32[512,129280]{1,0:T(8,128)}, /*index=80*/f32[256,1]{0,1:T(1,128)}, f32[512,1,256]{2,1,0:T(1,128)}, f32[256,1,512,512]{3,2,1,0:T(8,128)}, f32[256,1,512,512]{3,2,1,0:T(8,128)}, f32[256,1,512,512]{3,2,1,0:T(8,128)}, /*index=85*/f32[512,1,512]{2,1,0:T(1,128)}, f32[512,1,512]{2,1,0:T(1,128)}, f32[512,1,512]{2,1,0:T(1,128)}, f32[512,1]{0,1:T(1,128)}, f32[512,1]{0,1:T(1,128)}, /*index=90*/f32[512,1]{0,1:T(1,128)}, f32[128,1,128,512]{3,2,1,0:T(8,128)}, f32[1536,1]{0,1:T(1,128)}, f32[512,1,576]{0,2,1:T(8,128)}, f32[512,1,128,256]{3,2,1,0:T(8,128)}, /*index=95*/f32[512,1,1536]{2,1,0:T(1,128)}, f32[1536,1,128,192]{2,3,1,0:T(8,128)}, f32[129280,512]{1,0:T(8,128)}, s32[]{:T(128)}, f32[]{:T(128)}, /*index=100*/f32[]{:T(128)}, f32[]{:T(128)}, f32[]{:T(128)}, f32[]{:T(128)}, f32[]{:T(128)}, /*index=105*/f32[]{:T(128)}, f32[]{:T(128)}, f32[]{:T(128)}, s32[]{:T(128)}, f32[]{:T(128)})}, allow_spmd_sharding_propagation_to_parameters={false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false}, allow_spmd_sharding_propagation_to_output={false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,true,true,true,true,true,true,true,true,true,true,true}, num_partitions=4 +HloModule jit_train_step, is_scheduled=true, input_output_alias={ {0}: (0, {}, may-alias), {1}: (1, {}, may-alias), {2}: (2, {}, may-alias), {3}: (3, {}, may-alias), {4}: (4, {}, may-alias), {5}: (5, {}, may-alias), {6}: (6, {}, may-alias), {7}: (7, {}, may-alias), {8}: (8, {}, may-alias), {9}: (9, {}, may-alias), {10}: (10, {}, may-alias), {11}: (11, {}, may-alias), {12}: (12, {}, may-alias), {13}: (13, {}, may-alias), {14}: (14, {}, may-alias), {15}: (15, {}, may-alias), {16}: (16, {}, may-alias), {17}: (17, {}, may-alias), {18}: (18, {}, may-alias), {19}: (19, {}, may-alias), {20}: (20, {}, may-alias), {21}: (21, {}, may-alias), {22}: (22, {}, may-alias), {23}: (23, {}, may-alias), {24}: (24, {}, may-alias), {25}: (25, {}, may-alias), {26}: (26, {}, may-alias), {27}: (27, {}, may-alias), {28}: (28, {}, may-alias), {29}: (29, {}, may-alias), {30}: (30, {}, may-alias), {31}: (31, {}, may-alias), {32}: (32, {}, may-alias), {33}: (33, {}, may-alias), {34}: (34, {}, may-alias), {35}: (35, {}, may-alias), {36}: (36, {}, may-alias), {37}: (37, {}, may-alias), {38}: (38, {}, may-alias), {39}: (39, {}, may-alias), {40}: (40, {}, may-alias), {41}: (41, {}, may-alias), {42}: (42, {}, may-alias), {43}: (43, {}, may-alias), {44}: (44, {}, may-alias), {45}: (45, {}, may-alias), {46}: (46, {}, may-alias), {47}: (47, {}, may-alias), {48}: (48, {}, may-alias), {49}: (49, {}, may-alias), {50}: (50, {}, may-alias), {51}: (51, {}, may-alias), {52}: (52, {}, may-alias), {53}: (53, {}, may-alias), {54}: (54, {}, may-alias), {55}: (55, {}, may-alias), {56}: (56, {}, may-alias), {57}: (57, {}, may-alias), {58}: (58, {}, may-alias), {59}: (59, {}, may-alias), {60}: (60, {}, may-alias), {61}: (61, {}, may-alias), {62}: (62, {}, may-alias), {63}: (63, {}, may-alias), {64}: (64, {}, may-alias), {65}: (65, {}, may-alias), {66}: (66, {}, may-alias), {67}: (67, {}, may-alias), {68}: (68, {}, may-alias), {69}: (69, {}, may-alias), {70}: (70, {}, may-alias), {71}: (71, {}, may-alias), {72}: (72, {}, may-alias), {73}: (73, {}, may-alias), {74}: (74, {}, may-alias), {75}: (75, {}, may-alias), {76}: (76, {}, may-alias), {77}: (77, {}, may-alias), {78}: (78, {}, may-alias), {79}: (79, {}, may-alias), {80}: (80, {}, may-alias), {81}: (81, {}, may-alias), {82}: (82, {}, may-alias), {83}: (83, {}, may-alias), {84}: (84, {}, may-alias), {85}: (85, {}, may-alias), {86}: (86, {}, may-alias), {87}: (87, {}, may-alias), {88}: (88, {}, may-alias), {89}: (89, {}, may-alias), {90}: (90, {}, may-alias), {91}: (91, {}, may-alias), {92}: (92, {}, may-alias), {93}: (93, {}, may-alias), {94}: (94, {}, may-alias), {95}: (95, {}, may-alias), {96}: (96, {}, may-alias), {97}: (97, {}, may-alias), {98}: (98, {}, may-alias), {99}: (99, {}, may-alias), {100}: (100, {}, may-alias), {101}: (101, {}, may-alias), {102}: (102, {}, may-alias), {103}: (103, {}, may-alias), {104}: (104, {}, may-alias), {105}: (105, {}, may-alias), {106}: (106, {}, may-alias), {107}: (107, {}, may-alias), {108}: (108, {}, may-alias), {109}: (109, {}, may-alias), {110}: (110, {}, may-alias), {111}: (111, {}, may-alias), {112}: (112, {}, may-alias), {113}: (113, {}, may-alias), {114}: (114, {}, may-alias), {115}: (115, {}, may-alias), {116}: (116, {}, may-alias), {117}: (117, {}, may-alias), {118}: (118, {}, may-alias), {119}: (119, {}, may-alias), {120}: (120, {}, may-alias), {121}: (121, {}, may-alias), {122}: (122, {}, may-alias), {123}: (123, {}, may-alias), {124}: (124, {}, may-alias), {125}: (125, {}, may-alias), {126}: (126, {}, may-alias), {127}: (127, {}, may-alias), {128}: (128, {}, may-alias), {129}: (129, {}, may-alias), {130}: (130, {}, may-alias), {131}: (131, {}, may-alias), {132}: (132, {}, may-alias), {133}: (133, {}, may-alias), {134}: (134, {}, may-alias), {135}: (135, {}, may-alias), {136}: (136, {}, may-alias), {137}: (137, {}, may-alias), {138}: (138, {}, may-alias), {139}: (139, {}, may-alias), {140}: (140, {}, may-alias), {141}: (141, {}, may-alias), {142}: (142, {}, may-alias), {143}: (143, {}, may-alias), {144}: (144, {}, may-alias), {145}: (145, {}, may-alias), {146}: (146, {}, may-alias) }, entry_computation_layout={(f32[512]{0:T(512)}, u32[3]{0:T(128)}, u32[3,4]{1,0:T(4,128)}, u32[3]{0:T(128)}, u32[3,4]{1,0:T(4,128)}, /*index=5*/u32[3]{0:T(128)}, u32[3,4]{1,0:T(4,128)}, u32[3]{0:T(128)}, u32[3,4]{1,0:T(4,128)}, u32[3]{0:T(128)}, /*index=10*/u32[3,4]{1,0:T(4,128)}, u32[3]{0:T(128)}, u32[3,4]{1,0:T(4,128)}, f32[128,3,18432]{2,0,1:T(8,128)}, f32[128,3,18432]{2,0,1:T(8,128)}, /*index=15*/f32[18432,3,128]{2,0,1:T(8,128)}, f32[512,3]{0,1:T(4,128)}, f32[512,3]{0,1:T(4,128)}, u32[3]{0:T(128)}, u32[3,4]{1,0:T(4,128)}, /*index=20*/u32[3]{0:T(128)}, u32[3,4]{1,0:T(4,128)}, u32[3]{0:T(128)}, u32[3,4]{1,0:T(4,128)}, f32[512,3]{0,1:T(4,128)}, /*index=25*/f32[128,3,128,128]{3,2,1,0:T(8,128)}, f32[1536,3]{0,1:T(4,128)}, f32[128,3,576]{0,2,1:T(8,128)}, f32[128,3,128,256]{3,2,1,0:T(8,128)}, f32[128,3,1536]{2,0,1:T(8,128)}, /*index=30*/f32[384,3,128,192]{2,3,1,0:T(8,128)}, u32[]{:T(128)}, u32[4]{0:T(128)}, u32[]{:T(128)}, u32[4]{0:T(128)}, /*index=35*/u32[]{:T(128)}, u32[4]{0:T(128)}, f32[128,129280]{1,0:T(8,128)}, f32[256,1]{0,1:T(1,128)}, f32[128,1,256]{2,1,0:T(1,128)}, /*index=40*/u32[1]{0:T(128)}, u32[1,4]{1,0:T(1,128)}, u32[1]{0:T(128)}, u32[1,4]{1,0:T(1,128)}, u32[1]{0:T(128)}, /*index=45*/u32[1,4]{1,0:T(1,128)}, f32[256,1,128,512]{3,2,1,0:T(8,128)}, f32[256,1,128,512]{3,2,1,0:T(8,128)}, f32[256,1,512,128]{3,2,1,0:T(8,128)}, u32[1]{0:T(128)}, /*index=50*/u32[1,4]{1,0:T(1,128)}, u32[1]{0:T(128)}, u32[1,4]{1,0:T(1,128)}, u32[1]{0:T(128)}, u32[1,4]{1,0:T(1,128)}, /*index=55*/f32[128,1,512]{2,1,0:T(1,128)}, f32[128,1,512]{2,1,0:T(1,128)}, f32[512,1,128]{2,1,0:T(1,128)}, u32[1]{0:T(128)}, u32[1,4]{1,0:T(1,128)}, /*index=60*/u32[1]{0:T(128)}, u32[1,4]{1,0:T(1,128)}, u32[1]{0:T(128)}, u32[1,4]{1,0:T(1,128)}, f32[512,1]{0,1:T(1,128)}, /*index=65*/f32[512,1]{0,1:T(1,128)}, f32[512,1]{0,1:T(1,128)}, f32[128,1,128,128]{3,2,1,0:T(8,128)}, f32[1536,1]{0,1:T(1,128)}, f32[128,1,576]{0,2,1:T(8,128)}, /*index=70*/f32[128,1,128,256]{3,2,1,0:T(8,128)}, f32[128,1,1536]{2,1,0:T(1,128)}, f32[384,1,128,192]{2,3,1,0:T(8,128)}, u32[]{:T(128)}, u32[4]{0:T(128)}, /*index=75*/u32[]{:T(128)}, u32[4]{0:T(128)}, u32[]{:T(128)}, u32[4]{0:T(128)}, f32[129280,128]{1,0:T(8,128)}, /*index=80*/s32[]{:T(128)}, f32[512]{0:T(512)}, f32[128,3,18432]{2,0,1:T(8,128)}, f32[128,3,18432]{2,0,1:T(8,128)}, f32[18432,3,128]{2,0,1:T(8,128)}, /*index=85*/f32[512,3]{0,1:T(4,128)}, f32[512,3]{0,1:T(4,128)}, f32[512,3]{0,1:T(4,128)}, f32[128,3,128,128]{3,2,1,0:T(8,128)}, f32[1536,3]{0,1:T(4,128)}, /*index=90*/f32[128,3,576]{0,2,1:T(8,128)}, f32[128,3,128,256]{3,2,1,0:T(8,128)}, f32[128,3,1536]{2,0,1:T(8,128)}, f32[384,3,128,192]{2,3,1,0:T(8,128)}, f32[128,129280]{1,0:T(8,128)}, /*index=95*/f32[256,1]{0,1:T(1,128)}, f32[128,1,256]{2,1,0:T(1,128)}, f32[256,1,128,512]{3,2,1,0:T(8,128)}, f32[256,1,128,512]{3,2,1,0:T(8,128)}, f32[256,1,512,128]{3,2,1,0:T(8,128)}, /*index=100*/f32[128,1,512]{2,1,0:T(1,128)}, f32[128,1,512]{2,1,0:T(1,128)}, f32[512,1,128]{2,1,0:T(1,128)}, f32[512,1]{0,1:T(1,128)}, f32[512,1]{0,1:T(1,128)}, /*index=105*/f32[512,1]{0,1:T(1,128)}, f32[128,1,128,128]{3,2,1,0:T(8,128)}, f32[1536,1]{0,1:T(1,128)}, f32[128,1,576]{0,2,1:T(8,128)}, f32[128,1,128,256]{3,2,1,0:T(8,128)}, /*index=110*/f32[128,1,1536]{2,1,0:T(1,128)}, f32[384,1,128,192]{2,3,1,0:T(8,128)}, f32[129280,128]{1,0:T(8,128)}, f32[512]{0:T(512)}, f32[128,3,18432]{2,0,1:T(8,128)}, /*index=115*/f32[128,3,18432]{2,0,1:T(8,128)}, f32[18432,3,128]{2,0,1:T(8,128)}, f32[512,3]{0,1:T(4,128)}, f32[512,3]{0,1:T(4,128)}, f32[512,3]{0,1:T(4,128)}, /*index=120*/f32[128,3,128,128]{3,2,1,0:T(8,128)}, f32[1536,3]{0,1:T(4,128)}, f32[128,3,576]{0,2,1:T(8,128)}, f32[128,3,128,256]{3,2,1,0:T(8,128)}, f32[128,3,1536]{2,0,1:T(8,128)}, /*index=125*/f32[384,3,128,192]{2,3,1,0:T(8,128)}, f32[128,129280]{1,0:T(8,128)}, f32[256,1]{0,1:T(1,128)}, f32[128,1,256]{2,1,0:T(1,128)}, f32[256,1,128,512]{3,2,1,0:T(8,128)}, /*index=130*/f32[256,1,128,512]{3,2,1,0:T(8,128)}, f32[256,1,512,128]{3,2,1,0:T(8,128)}, f32[128,1,512]{2,1,0:T(1,128)}, f32[128,1,512]{2,1,0:T(1,128)}, f32[512,1,128]{2,1,0:T(1,128)}, /*index=135*/f32[512,1]{0,1:T(1,128)}, f32[512,1]{0,1:T(1,128)}, f32[512,1]{0,1:T(1,128)}, f32[128,1,128,128]{3,2,1,0:T(8,128)}, f32[1536,1]{0,1:T(1,128)}, /*index=140*/f32[128,1,576]{0,2,1:T(8,128)}, f32[128,1,128,256]{3,2,1,0:T(8,128)}, f32[128,1,1536]{2,1,0:T(1,128)}, f32[384,1,128,192]{2,3,1,0:T(8,128)}, f32[129280,128]{1,0:T(8,128)}, /*index=145*/s32[]{:T(128)}, u32[]{:T(128)}, s32[1,128]{1,0:T(1,128)}, s32[1,128]{1,0:T(1,128)}, s32[1,128]{1,0:T(1,128)}, /*index=150*/s32[1,128]{1,0:T(1,128)}, s32[1,128]{1,0:T(1,128)})->(f32[512]{0:T(512)}, u32[3]{0:T(128)}, u32[3,4]{1,0:T(4,128)}, u32[3]{0:T(128)}, u32[3,4]{1,0:T(4,128)}, /*index=5*/u32[3]{0:T(128)}, u32[3,4]{1,0:T(4,128)}, u32[3]{0:T(128)}, u32[3,4]{1,0:T(4,128)}, u32[3]{0:T(128)}, /*index=10*/u32[3,4]{1,0:T(4,128)}, u32[3]{0:T(128)}, u32[3,4]{1,0:T(4,128)}, f32[128,3,18432]{2,0,1:T(8,128)}, f32[128,3,18432]{2,0,1:T(8,128)}, /*index=15*/f32[18432,3,128]{2,0,1:T(8,128)}, f32[512,3]{0,1:T(4,128)}, f32[512,3]{0,1:T(4,128)}, u32[3]{0:T(128)}, u32[3,4]{1,0:T(4,128)}, /*index=20*/u32[3]{0:T(128)}, u32[3,4]{1,0:T(4,128)}, u32[3]{0:T(128)}, u32[3,4]{1,0:T(4,128)}, f32[512,3]{0,1:T(4,128)}, /*index=25*/f32[128,3,128,128]{3,2,1,0:T(8,128)}, f32[1536,3]{0,1:T(4,128)}, f32[128,3,576]{0,2,1:T(8,128)}, f32[128,3,128,256]{3,2,1,0:T(8,128)}, f32[128,3,1536]{2,0,1:T(8,128)}, /*index=30*/f32[384,3,128,192]{2,3,1,0:T(8,128)}, u32[]{:T(128)}, u32[4]{0:T(128)}, u32[]{:T(128)}, u32[4]{0:T(128)}, /*index=35*/u32[]{:T(128)}, u32[4]{0:T(128)}, f32[128,129280]{1,0:T(8,128)}, f32[256,1]{0,1:T(1,128)}, f32[128,1,256]{2,1,0:T(1,128)}, /*index=40*/u32[1]{0:T(128)}, u32[1,4]{1,0:T(1,128)}, u32[1]{0:T(128)}, u32[1,4]{1,0:T(1,128)}, u32[1]{0:T(128)}, /*index=45*/u32[1,4]{1,0:T(1,128)}, f32[256,1,128,512]{3,2,1,0:T(8,128)}, f32[256,1,128,512]{3,2,1,0:T(8,128)}, f32[256,1,512,128]{3,2,1,0:T(8,128)}, u32[1]{0:T(128)}, /*index=50*/u32[1,4]{1,0:T(1,128)}, u32[1]{0:T(128)}, u32[1,4]{1,0:T(1,128)}, u32[1]{0:T(128)}, u32[1,4]{1,0:T(1,128)}, /*index=55*/f32[128,1,512]{2,1,0:T(1,128)}, f32[128,1,512]{2,1,0:T(1,128)}, f32[512,1,128]{2,1,0:T(1,128)}, u32[1]{0:T(128)}, u32[1,4]{1,0:T(1,128)}, /*index=60*/u32[1]{0:T(128)}, u32[1,4]{1,0:T(1,128)}, u32[1]{0:T(128)}, u32[1,4]{1,0:T(1,128)}, f32[512,1]{0,1:T(1,128)}, /*index=65*/f32[512,1]{0,1:T(1,128)}, f32[512,1]{0,1:T(1,128)}, f32[128,1,128,128]{3,2,1,0:T(8,128)}, f32[1536,1]{0,1:T(1,128)}, f32[128,1,576]{0,2,1:T(8,128)}, /*index=70*/f32[128,1,128,256]{3,2,1,0:T(8,128)}, f32[128,1,1536]{2,1,0:T(1,128)}, f32[384,1,128,192]{2,3,1,0:T(8,128)}, u32[]{:T(128)}, u32[4]{0:T(128)}, /*index=75*/u32[]{:T(128)}, u32[4]{0:T(128)}, u32[]{:T(128)}, u32[4]{0:T(128)}, f32[129280,128]{1,0:T(8,128)}, /*index=80*/s32[]{:T(128)}, f32[512]{0:T(512)}, f32[128,3,18432]{2,0,1:T(8,128)}, f32[128,3,18432]{2,0,1:T(8,128)}, f32[18432,3,128]{2,0,1:T(8,128)}, /*index=85*/f32[512,3]{0,1:T(4,128)}, f32[512,3]{0,1:T(4,128)}, f32[512,3]{0,1:T(4,128)}, f32[128,3,128,128]{3,2,1,0:T(8,128)}, f32[1536,3]{0,1:T(4,128)}, /*index=90*/f32[128,3,576]{0,2,1:T(8,128)}, f32[128,3,128,256]{3,2,1,0:T(8,128)}, f32[128,3,1536]{2,0,1:T(8,128)}, f32[384,3,128,192]{2,3,1,0:T(8,128)}, f32[128,129280]{1,0:T(8,128)}, /*index=95*/f32[256,1]{0,1:T(1,128)}, f32[128,1,256]{2,1,0:T(1,128)}, f32[256,1,128,512]{3,2,1,0:T(8,128)}, f32[256,1,128,512]{3,2,1,0:T(8,128)}, f32[256,1,512,128]{3,2,1,0:T(8,128)}, /*index=100*/f32[128,1,512]{2,1,0:T(1,128)}, f32[128,1,512]{2,1,0:T(1,128)}, f32[512,1,128]{2,1,0:T(1,128)}, f32[512,1]{0,1:T(1,128)}, f32[512,1]{0,1:T(1,128)}, /*index=105*/f32[512,1]{0,1:T(1,128)}, f32[128,1,128,128]{3,2,1,0:T(8,128)}, f32[1536,1]{0,1:T(1,128)}, f32[128,1,576]{0,2,1:T(8,128)}, f32[128,1,128,256]{3,2,1,0:T(8,128)}, /*index=110*/f32[128,1,1536]{2,1,0:T(1,128)}, f32[384,1,128,192]{2,3,1,0:T(8,128)}, f32[129280,128]{1,0:T(8,128)}, f32[512]{0:T(512)}, f32[128,3,18432]{2,0,1:T(8,128)}, /*index=115*/f32[128,3,18432]{2,0,1:T(8,128)}, f32[18432,3,128]{2,0,1:T(8,128)}, f32[512,3]{0,1:T(4,128)}, f32[512,3]{0,1:T(4,128)}, f32[512,3]{0,1:T(4,128)}, /*index=120*/f32[128,3,128,128]{3,2,1,0:T(8,128)}, f32[1536,3]{0,1:T(4,128)}, f32[128,3,576]{0,2,1:T(8,128)}, f32[128,3,128,256]{3,2,1,0:T(8,128)}, f32[128,3,1536]{2,0,1:T(8,128)}, /*index=125*/f32[384,3,128,192]{2,3,1,0:T(8,128)}, f32[128,129280]{1,0:T(8,128)}, f32[256,1]{0,1:T(1,128)}, f32[128,1,256]{2,1,0:T(1,128)}, f32[256,1,128,512]{3,2,1,0:T(8,128)}, /*index=130*/f32[256,1,128,512]{3,2,1,0:T(8,128)}, f32[256,1,512,128]{3,2,1,0:T(8,128)}, f32[128,1,512]{2,1,0:T(1,128)}, f32[128,1,512]{2,1,0:T(1,128)}, f32[512,1,128]{2,1,0:T(1,128)}, /*index=135*/f32[512,1]{0,1:T(1,128)}, f32[512,1]{0,1:T(1,128)}, f32[512,1]{0,1:T(1,128)}, f32[128,1,128,128]{3,2,1,0:T(8,128)}, f32[1536,1]{0,1:T(1,128)}, /*index=140*/f32[128,1,576]{0,2,1:T(8,128)}, f32[128,1,128,256]{3,2,1,0:T(8,128)}, f32[128,1,1536]{2,1,0:T(1,128)}, f32[384,1,128,192]{2,3,1,0:T(8,128)}, f32[129280,128]{1,0:T(8,128)}, /*index=145*/s32[]{:T(128)}, u32[]{:T(128)}, f32[]{:T(128)}, f32[]{:T(128)}, f32[]{:T(128)}, /*index=150*/f32[]{:T(128)}, f32[]{:T(128)}, f32[]{:T(128)}, f32[]{:T(128)}, f32[]{:T(128)}, /*index=155*/f32[]{:T(128)}, s32[]{:T(128)}, f32[]{:T(128)})}, allow_spmd_sharding_propagation_to_parameters={false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false}, allow_spmd_sharding_propagation_to_output={false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,true,true,true,true,true,true,true,true,true,true,true}, num_partitions=4, frontend_attributes={arg_layout_modes="default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default",arg_memory_spaces="0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0",out_layout_modes="default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default",out_memory_spaces="0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0"} FileNames @@ -9,1992 +9,1992 @@ FileLocations StackFrames -%region_46.56 (top_k.25: bf16[], top_k.26: bf16[], top_k.27: s32[], top_k.28: s32[]) -> pred[] { - %constant.1427 = s32[]{:T(128)} constant(0) - %constant.1428 = s32[]{:T(128)} constant(2147483647) - %top_k.25 = bf16[]{:T(256)} parameter(0), metadata={op_name="top_k"} - %top_k.26 = bf16[]{:T(256)} parameter(1), metadata={op_name="top_k"} - %top_k.27 = s32[]{:T(128)} parameter(2), metadata={op_name="top_k"} - %top_k.28 = s32[]{:T(128)} parameter(3), metadata={op_name="top_k"} - %convert.393 = f32[]{:T(128)S(6)} convert(%top_k.25), metadata={op_name="convert.18"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} - %bitcast-convert.39 = s32[]{:T(128)S(6)} bitcast-convert(%convert.393), metadata={op_name="bitcast-convert.8"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} - %compare.144 = pred[]{:T(512)S(6)} compare(%bitcast-convert.39, %constant.1427), direction=LT, metadata={op_name="compare.38"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} - %xor.40 = s32[]{:T(128)S(6)} xor(%constant.1428, %bitcast-convert.39), metadata={op_name="xor.8"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} - %select.127 = s32[]{:T(128)S(6)} select(%compare.144, %xor.40, %bitcast-convert.39), metadata={op_name="select.16"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[{"indices":["1","3"]}]}} - %convert.394 = f32[]{:T(128)S(6)} convert(%top_k.26), metadata={op_name="convert.19"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} - %bitcast-convert.40 = s32[]{:T(128)S(6)} bitcast-convert(%convert.394), metadata={op_name="bitcast-convert.9"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} - %compare.145 = pred[]{:T(512)S(6)} compare(%bitcast-convert.40, %constant.1427), direction=LT, metadata={op_name="compare.39"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} - %xor.41 = s32[]{:T(128)S(6)} xor(%constant.1428, %bitcast-convert.40), metadata={op_name="xor.9"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} - %select.128 = s32[]{:T(128)S(6)} select(%compare.145, %xor.41, %bitcast-convert.40), metadata={op_name="select.17"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[{"indices":["1","3"]}]}} - %compare.146 = pred[]{:T(512)S(6)} compare(%select.127, %select.128), direction=GT, metadata={op_name="compare.0"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} - %compare.147 = pred[]{:T(512)S(6)} compare(%select.128, %select.127), direction=GT, metadata={op_name="compare.117"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} - %compare.148 = pred[]{:T(512)S(6)} compare(%compare.146, %compare.147), direction=EQ, metadata={op_name="compare.118"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} - %compare.149 = pred[]{:T(512)S(6)} compare(%top_k.27, %top_k.28), direction=LT, metadata={op_name="compare.119"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} - ROOT %select.129 = pred[]{:T(512)} select(%compare.148, %compare.149, %compare.146), metadata={op_name="select.113"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_47.57 (sort.64: s32[], sort.65: s32[], sort.66: s32[], sort.67: s32[]) -> pred[] { - %sort.64 = s32[]{:T(128)} parameter(0), metadata={op_name="jit(argsort)/sort"} - %sort.65 = s32[]{:T(128)} parameter(1), metadata={op_name="jit(argsort)/sort"} - %sort.66 = s32[]{:T(128)} parameter(2), metadata={op_name="jit(argsort)/sort"} - %sort.67 = s32[]{:T(128)} parameter(3), metadata={op_name="jit(argsort)/sort"} - %lt_to.32 = pred[]{:T(512)S(6)} compare(%sort.64, %sort.65), direction=LT, metadata={op_name="lt_to"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} - %lt_to.33 = pred[]{:T(512)S(6)} compare(%sort.65, %sort.64), direction=LT, metadata={op_name="lt_to"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} - %compare.150 = pred[]{:T(512)S(6)} compare(%lt_to.32, %lt_to.33), direction=EQ, metadata={op_name="compare.120"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} - %compare.151 = pred[]{:T(512)S(6)} compare(%sort.66, %sort.67), direction=LT, metadata={op_name="compare.121"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} - ROOT %select.130 = pred[]{:T(512)} select(%compare.150, %compare.151, %lt_to.32), metadata={op_name="select.114"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_48.58 (sort.68: s32[], sort.69: s32[], sort.70: s32[], sort.71: s32[]) -> pred[] { - %sort.68 = s32[]{:T(128)} parameter(0), metadata={op_name="jit(argsort)/sort"} - %sort.69 = s32[]{:T(128)} parameter(1), metadata={op_name="jit(argsort)/sort"} - %sort.70 = s32[]{:T(128)} parameter(2), metadata={op_name="jit(argsort)/sort"} - %sort.71 = s32[]{:T(128)} parameter(3), metadata={op_name="jit(argsort)/sort"} - %lt_to.34 = pred[]{:T(512)S(6)} compare(%sort.68, %sort.69), direction=LT, metadata={op_name="lt_to"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} - %lt_to.35 = pred[]{:T(512)S(6)} compare(%sort.69, %sort.68), direction=LT, metadata={op_name="lt_to"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} - %compare.152 = pred[]{:T(512)S(6)} compare(%lt_to.34, %lt_to.35), direction=EQ, metadata={op_name="compare.122"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} - %compare.153 = pred[]{:T(512)S(6)} compare(%sort.70, %sort.71), direction=LT, metadata={op_name="compare.123"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} - ROOT %select.131 = pred[]{:T(512)} select(%compare.152, %compare.153, %lt_to.34), metadata={op_name="select.115"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_67.80 (sort.78: s32[], sort.79: s32[], sort.80: s32[], sort.81: s32[]) -> pred[] { - %sort.78 = s32[]{:T(128)} parameter(0), metadata={op_name="sort_activations/jit(argsort)/sort"} - %sort.79 = s32[]{:T(128)} parameter(1), metadata={op_name="sort_activations/jit(argsort)/sort"} - %sort.80 = s32[]{:T(128)} parameter(2), metadata={op_name="sort_activations/jit(argsort)/sort"} - %sort.81 = s32[]{:T(128)} parameter(3), metadata={op_name="sort_activations/jit(argsort)/sort"} - %lt_to.37 = pred[]{:T(512)S(6)} compare(%sort.78, %sort.79), direction=LT, metadata={op_name="lt_to"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} - %lt_to.38 = pred[]{:T(512)S(6)} compare(%sort.79, %sort.78), direction=LT, metadata={op_name="lt_to"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} - %compare.156 = pred[]{:T(512)S(6)} compare(%lt_to.37, %lt_to.38), direction=EQ, metadata={op_name="compare.124"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} - %compare.157 = pred[]{:T(512)S(6)} compare(%sort.80, %sort.81), direction=LT, metadata={op_name="compare.125"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} - ROOT %select.134 = pred[]{:T(512)} select(%compare.156, %compare.157, %lt_to.37), metadata={op_name="select.116"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_119.141 (reduce_sum.157: bf16[], reduce_sum.158: bf16[]) -> bf16[] { - %reduce_sum.157 = bf16[]{:T(256)} parameter(0), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/reduce_sum"} - %reduce_sum.158 = bf16[]{:T(256)} parameter(1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/reduce_sum"} - ROOT %reduce_sum.159 = bf16[]{:T(256)} add(%reduce_sum.157, %reduce_sum.158), metadata={op_name="checkpoint/moe_layers/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_107.126 (psum.6: bf16[], psum.9: bf16[]) -> bf16[] { - %psum.6 = bf16[]{:T(256)} parameter(0), metadata={op_name="psum"} - %psum.9 = bf16[]{:T(256)} parameter(1), metadata={op_name="psum"} - ROOT %add.1445 = bf16[]{:T(256)} add(%psum.6, %psum.9), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_108.127 (psum.10: bf16[], psum.11: bf16[]) -> bf16[] { - %psum.10 = bf16[]{:T(256)} parameter(0), metadata={op_name="psum"} - %psum.11 = bf16[]{:T(256)} parameter(1), metadata={op_name="psum"} - ROOT %add.1446 = bf16[]{:T(256)} add(%psum.10, %psum.11), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_109.128 (psum.14: bf16[], psum.15: bf16[]) -> bf16[] { - %psum.14 = bf16[]{:T(256)} parameter(0), metadata={op_name="psum"} - %psum.15 = bf16[]{:T(256)} parameter(1), metadata={op_name="psum"} - ROOT %add.1447 = bf16[]{:T(256)} add(%psum.14, %psum.15), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_62.73 (reduce-window.111: s32[], reduce-window.112: s32[]) -> s32[] { - %reduce-window.111 = s32[]{:T(128)} parameter(0), metadata={op_name="reduce-window.35"} - %reduce-window.112 = s32[]{:T(128)} parameter(1), metadata={op_name="reduce-window.35"} - ROOT %reduce_window_sum.108 = s32[]{:T(128)} add(%reduce-window.111, %reduce-window.112), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_64.75 (reduce-window.113: s32[], reduce-window.114: s32[]) -> s32[] { - %reduce-window.113 = s32[]{:T(128)} parameter(0), metadata={op_name="reduce-window.36"} - %reduce-window.114 = s32[]{:T(128)} parameter(1), metadata={op_name="reduce-window.36"} - ROOT %reduce_window_sum.109 = s32[]{:T(128)} add(%reduce-window.113, %reduce-window.114), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_65.76 (reduce-window.115: s32[], reduce-window.116: s32[]) -> s32[] { - %reduce-window.115 = s32[]{:T(128)} parameter(0), metadata={op_name="reduce-window.63"} - %reduce-window.116 = s32[]{:T(128)} parameter(1), metadata={op_name="reduce-window.63"} - ROOT %reduce_window_sum.110 = s32[]{:T(128)} add(%reduce-window.115, %reduce-window.116), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_68.81.clone (reduce-window.396: s32[], reduce-window.397: s32[]) -> s32[] { - %reduce-window.396 = s32[]{:T(128)} parameter(0), metadata={op_name="reduce-window.38"} - %reduce-window.397 = s32[]{:T(128)} parameter(1), metadata={op_name="reduce-window.38"} - ROOT %reduce_window_sum.317 = s32[]{:T(128)} add(%reduce-window.396, %reduce-window.397), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_69.82.clone (reduce-window.400: s32[], reduce-window.401: s32[]) -> s32[] { - %reduce-window.400 = s32[]{:T(128)} parameter(0), metadata={op_name="reduce-window.64"} - %reduce-window.401 = s32[]{:T(128)} parameter(1), metadata={op_name="reduce-window.64"} - ROOT %reduce_window_sum.319 = s32[]{:T(128)} add(%reduce-window.400, %reduce-window.401), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_71.84.clone (reduce-window.404: s32[], reduce-window.405: s32[]) -> s32[] { - %reduce-window.404 = s32[]{:T(128)} parameter(0), metadata={op_name="reduce-window.40"} - %reduce-window.405 = s32[]{:T(128)} parameter(1), metadata={op_name="reduce-window.40"} - ROOT %reduce_window_sum.321 = s32[]{:T(128)} add(%reduce-window.404, %reduce-window.405), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_72.85.clone (reduce-window.408: s32[], reduce-window.409: s32[]) -> s32[] { - %reduce-window.408 = s32[]{:T(128)} parameter(0), metadata={op_name="reduce-window.65"} - %reduce-window.409 = s32[]{:T(128)} parameter(1), metadata={op_name="reduce-window.65"} - ROOT %reduce_window_sum.323 = s32[]{:T(128)} add(%reduce-window.408, %reduce-window.409), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_74.87.clone (reduce-window.412: s32[], reduce-window.413: s32[]) -> s32[] { - %reduce-window.412 = s32[]{:T(128)} parameter(0), metadata={op_name="reduce-window.42"} - %reduce-window.413 = s32[]{:T(128)} parameter(1), metadata={op_name="reduce-window.42"} - ROOT %reduce_window_sum.325 = s32[]{:T(128)} add(%reduce-window.412, %reduce-window.413), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_75.88.clone (reduce-window.416: s32[], reduce-window.417: s32[]) -> s32[] { - %reduce-window.416 = s32[]{:T(128)} parameter(0), metadata={op_name="reduce-window.66"} - %reduce-window.417 = s32[]{:T(128)} parameter(1), metadata={op_name="reduce-window.66"} - ROOT %reduce_window_sum.327 = s32[]{:T(128)} add(%reduce-window.416, %reduce-window.417), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_80.96.clone (reduce-window.420: s32[], reduce-window.421: s32[]) -> s32[] { - %reduce-window.420 = s32[]{:T(128)} parameter(0), metadata={op_name="reduce-window.44"} - %reduce-window.421 = s32[]{:T(128)} parameter(1), metadata={op_name="reduce-window.44"} - ROOT %reduce_window_sum.329 = s32[]{:T(128)} add(%reduce-window.420, %reduce-window.421), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_82.98.clone (reduce-window.424: s32[], reduce-window.425: s32[]) -> s32[] { - %reduce-window.424 = s32[]{:T(128)} parameter(0), metadata={op_name="reduce-window.45"} - %reduce-window.425 = s32[]{:T(128)} parameter(1), metadata={op_name="reduce-window.45"} - ROOT %reduce_window_sum.331 = s32[]{:T(128)} add(%reduce-window.424, %reduce-window.425), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_83.99.clone (reduce-window.428: s32[], reduce-window.429: s32[]) -> s32[] { - %reduce-window.428 = s32[]{:T(128)} parameter(0), metadata={op_name="reduce-window.67"} - %reduce-window.429 = s32[]{:T(128)} parameter(1), metadata={op_name="reduce-window.67"} - ROOT %reduce_window_sum.333 = s32[]{:T(128)} add(%reduce-window.428, %reduce-window.429), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_94.112 (reduce-window.174: s32[], reduce-window.175: s32[]) -> s32[] { - %reduce-window.174 = s32[]{:T(128)} parameter(0), metadata={op_name="reduce-window.49"} - %reduce-window.175 = s32[]{:T(128)} parameter(1), metadata={op_name="reduce-window.49"} - ROOT %reduce_window_sum.138 = s32[]{:T(128)} add(%reduce-window.174, %reduce-window.175), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_95.113 (reduce-window.179: s32[], reduce-window.180: s32[]) -> s32[] { - %reduce-window.179 = s32[]{:T(128)} parameter(0), metadata={op_name="reduce-window.69"} - %reduce-window.180 = s32[]{:T(128)} parameter(1), metadata={op_name="reduce-window.69"} - ROOT %reduce_window_sum.139 = s32[]{:T(128)} add(%reduce-window.179, %reduce-window.180), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_97.115 (reduce-window.184: s32[], reduce-window.185: s32[]) -> s32[] { - %reduce-window.184 = s32[]{:T(128)} parameter(0), metadata={op_name="reduce-window.51"} - %reduce-window.185 = s32[]{:T(128)} parameter(1), metadata={op_name="reduce-window.51"} - ROOT %reduce_window_sum.140 = s32[]{:T(128)} add(%reduce-window.184, %reduce-window.185), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_98.116 (reduce-window.189: s32[], reduce-window.190: s32[]) -> s32[] { - %reduce-window.189 = s32[]{:T(128)} parameter(0), metadata={op_name="reduce-window.70"} - %reduce-window.190 = s32[]{:T(128)} parameter(1), metadata={op_name="reduce-window.70"} - ROOT %reduce_window_sum.141 = s32[]{:T(128)} add(%reduce-window.189, %reduce-window.190), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_103.121 (reduce-window.194: s32[], reduce-window.195: s32[]) -> s32[] { - %reduce-window.194 = s32[]{:T(128)} parameter(0), metadata={op_name="reduce-window.53"} - %reduce-window.195 = s32[]{:T(128)} parameter(1), metadata={op_name="reduce-window.53"} - ROOT %reduce_window_sum.142 = s32[]{:T(128)} add(%reduce-window.194, %reduce-window.195), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_105.123 (reduce-window.199: s32[], reduce-window.200: s32[]) -> s32[] { - %reduce-window.199 = s32[]{:T(128)} parameter(0), metadata={op_name="reduce-window.54"} - %reduce-window.200 = s32[]{:T(128)} parameter(1), metadata={op_name="reduce-window.54"} - ROOT %reduce_window_sum.143 = s32[]{:T(128)} add(%reduce-window.199, %reduce-window.200), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +%fused_computation.1355 (param_0.3835: s32[1,128]) -> s32[1,1,128] { + %param_0.3835 = s32[1,128]{1,0:T(1,128)S(1)} parameter(0) + %constant.1719.clone.154 = s32[]{:T(128)} constant(0), metadata={op_name="jit(train_step)/shard_map/jit(gmm)" stack_frame_id=0} + %broadcast.4423 = s32[1,128]{1,0:T(1,128)} broadcast(%constant.1719.clone.154), dimensions={}, metadata={op_name="broadcast.370"} + %lt.832 = pred[1,128]{1,0:T(4,128)(4,1)} compare(%param_0.3835, %broadcast.4423), direction=LT, metadata={op_name="jit(train_step)/jvp()/lt" stack_frame_id=0} + %constant.1731.clone.1 = s32[]{:T(128)} constant(129280) + %add.3653 = s32[1,128]{1,0:T(1,128)} broadcast(%constant.1731.clone.1), dimensions={}, metadata={op_name="jit(train_step)/jvp()/add" stack_frame_id=0} + %add.3589 = s32[1,128]{1,0:T(1,128)} add(%param_0.3835, %add.3653), metadata={op_name="jit(train_step)/jvp()/add" stack_frame_id=0} + %select_n.2272 = s32[1,128]{1,0:T(1,128)} select(%lt.832, %add.3589, %param_0.3835), metadata={op_name="jit(train_step)/jvp()/select_n" stack_frame_id=0} + ROOT %bitcast.1928 = s32[1,1,128]{2,1,0:T(1,128)S(1)} bitcast(%select_n.2272) +} + +%fused_computation.1133 (param_0.3323: s32[512]) -> s32[1024] { + %constant.4122 = s32[] constant(0), metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} + %broadcast.4269 = s32[1024]{0:T(1024)} broadcast(%constant.4122), dimensions={}, metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} + %param_0.3323 = s32[512]{0:T(512)S(1)} parameter(0) + %constant.4128 = s32[] constant(2147483647), metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} + %pad.332 = s32[1024]{0:T(1024)} pad(%param_0.3323, %constant.4128), padding=0_512, metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} + %constant.4111 = s32[] constant(129279), metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} + %broadcast.4258 = s32[1024]{0:T(1024)} broadcast(%constant.4111), dimensions={}, metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} + ROOT %clamp.60 = s32[1024]{0:T(1024)S(1)} clamp(%broadcast.4269, %pad.332, %broadcast.4258), metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} +} + +%fused_computation.6 (param_0.20: bf16[129280,128], param_1.121: s32[1024]) -> bf16[512,128] { + %param_0.20 = bf16[129280,128]{1,0:T(8,128)(2,1)S(1)} parameter(0) + %param_1.121 = s32[1024]{0:T(1024)S(1)} parameter(1) + %custom-call.18 = s32[1024]{0:T(1024)} custom-call(%param_1.121), custom_call_target="AssumeGatherIndicesInBound", operand_layout_constraints={s32[1024]{0:T(1024)}}, metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} + %slice.1057 = s32[512]{0:T(512)} slice(%custom-call.18), slice={[0:512]}, metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} + %reshape.4326 = s32[4,128]{1,0:T(4,128)} reshape(%slice.1057), metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} + %transpose.891 = s32[4,128]{1,0:T(4,128)} transpose(%reshape.4326), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} + %gather.226 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)} gather(%param_0.20, %transpose.891), offset_dims={2}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=2, slice_sizes={1,128}, metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} + %transpose.890 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)} transpose(%gather.226), dimensions={0,1,2}, metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} + ROOT %reshape.4325 = bf16[512,128]{1,0:T(8,128)(2,1)S(1)} reshape(%transpose.890), metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} +} + +%fused_computation.1432 (param_0.4405: f32[512,3]) -> bf16[3,512] { + %param_0.4405 = f32[512,3]{0,1:T(4,128)S(1)} parameter(0) + %bitcast.2157 = f32[3,512]{1,0:T(4,128)} bitcast(%param_0.4405), metadata={op_name="jit(train_step)/jvp()/transpose" stack_frame_id=0} + ROOT %convert.640 = bf16[3,512]{1,0:T(4,128)(2,1)} convert(%bitcast.2157) +} + +%fused_computation.1122 (param_0.3211: f32[1536,3], param_1.3571: s32[]) -> bf16[3,384] { + %param_0.3211 = f32[1536,3]{0,1:T(4,128)S(1)} parameter(0) + %param_1.3571 = s32[]{:T(128)S(6)} parameter(1) + %constant.1719.clone.47 = s32[]{:T(128)} constant(0), metadata={op_name="jit(train_step)/shard_map/jit(gmm)" stack_frame_id=0} + %dynamic-slice.424 = bf16[384,3]{0,1:T(4,128)(2,1)} dynamic-slice(%param_0.3211, %param_1.3571, %constant.1719.clone.47), dynamic_slice_sizes={384,3}, metadata={op_name="jit(train_step)/jvp()/transpose" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294965375","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + ROOT %bitcast.1876 = bf16[3,384]{1,0:T(4,128)(2,1)} bitcast(%dynamic-slice.424), metadata={op_name="jit(train_step)/jvp()/transpose" stack_frame_id=0} +} + +%fused_computation.1431 (param_0.4404: f32[512,3]) -> bf16[3,512] { + %param_0.4404 = f32[512,3]{0,1:T(4,128)S(1)} parameter(0) + %bitcast.2156 = f32[3,512]{1,0:T(4,128)} bitcast(%param_0.4404), metadata={op_name="jit(train_step)/jvp()/transpose" stack_frame_id=0} + ROOT %convert.638 = bf16[3,512]{1,0:T(4,128)(2,1)} convert(%bitcast.2156) +} + +%fused_computation.1425 (param_0.4386: f32[128,3,1536]) -> bf16[3,128,1536] { + %param_0.4386 = f32[128,3,1536]{2,0,1:T(8,128)S(1)} parameter(0) + %bitcast.2155 = f32[3,128,1536]{2,1,0:T(8,128)} bitcast(%param_0.4386), metadata={op_name="jit(train_step)/jvp()/transpose" stack_frame_id=0} + ROOT %convert.636 = bf16[3,128,1536]{2,1,0:T(8,128)(2,1)} convert(%bitcast.2155) +} + +%fused_computation.484 (param_0.1460: f32[384,3,128,192]) -> bf16[3,384,128,192] { + %param_0.1460 = f32[384,3,128,192]{2,3,1,0:T(8,128)} parameter(0) + %copy.1445 = bf16[384,3,128,192]{2,0,3,1:T(8,128)(2,1)} copy(%param_0.1460), sharding={devices=[4,1,1,1]<=[4]}, metadata={op_name="state[\'model\'][\'decoder\'][\'dense_layers\'][\'self_attention\'][\'wq_b\'][\'kernel\'].value"} + ROOT %bitcast.1106 = bf16[3,384,128,192]{2,1,3,0:T(8,128)(2,1)} bitcast(%copy.1445), metadata={op_name="jit(train_step)/jvp()/transpose" stack_frame_id=0} +} + +%fused_computation.1256 (param_0.3509: f32[512,3], param_1.3939: s32[]) -> bf16[3,128] { + %param_0.3509 = f32[512,3]{0,1:T(4,128)S(1)} parameter(0) + %param_1.3939 = s32[]{:T(128)S(6)} parameter(1) + %constant.1719.clone.59 = s32[]{:T(128)} constant(0), metadata={op_name="jit(train_step)/shard_map/jit(gmm)" stack_frame_id=0} + %dynamic-slice.442 = bf16[128,3]{0,1:T(4,128)(2,1)} dynamic-slice(%param_0.3509, %param_1.3939, %constant.1719.clone.59), dynamic_slice_sizes={128,3}, metadata={op_name="jit(train_step)/jvp()/transpose" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294966911","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + ROOT %bitcast.1907 = bf16[3,128]{1,0:T(4,128)(2,1)} bitcast(%dynamic-slice.442), metadata={op_name="jit(train_step)/jvp()/transpose" stack_frame_id=0} +} + +%fused_computation.569 (param_0.1678: f32[128,3,128,256]) -> bf16[3,128,128,256] { + %param_0.1678 = f32[128,3,128,256]{3,2,1,0:T(8,128)S(1)} parameter(0) + %copy.1446 = bf16[128,3,128,256]{3,0,2,1:T(8,128)(2,1)} copy(%param_0.1678), sharding={devices=[4,1,1,1]<=[4]}, metadata={op_name="state[\'model\'][\'decoder\'][\'dense_layers\'][\'self_attention\'][\'wkv_b\'][\'kernel\'].value"} + ROOT %bitcast.1207 = bf16[3,128,128,256]{3,1,2,0:T(8,128)(2,1)} bitcast(%copy.1446), metadata={op_name="jit(train_step)/jvp()/transpose" stack_frame_id=0} +} + +%fused_computation.1395 () -> f32[32] { + %constant.1736.clone.1 = f32[]{:T(128)} constant(10000) + %broadcast.4415 = f32[32]{0:T(128)} broadcast(%constant.1736.clone.1), dimensions={}, metadata={op_name="broadcast.1660"} + %iota.346 = f32[32]{0:T(128)} iota(), iota_dimension=0, metadata={op_name="jit(train_step)/iota" stack_frame_id=0} + %constant.1737.clone.1 = f32[]{:T(128)} constant(0.03125) + %div.2470 = f32[32]{0:T(128)} broadcast(%constant.1737.clone.1), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.2443 = f32[32]{0:T(128)} multiply(%iota.346, %div.2470), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %pow.47 = f32[32]{0:T(128)} negate(%div.2443), metadata={op_name="jit(train_step)/pow" stack_frame_id=0} + %pow.46 = f32[32]{0:T(128)} power(%broadcast.4415, %pow.47), metadata={op_name="jit(train_step)/pow" stack_frame_id=0} + %constant.1738.clone.1 = f32[]{:T(128)} constant(0.025) + %broadcast.4414 = f32[32]{0:T(128)} broadcast(%constant.1738.clone.1), dimensions={}, metadata={op_name="broadcast.1662"} + %div.2442 = f32[32]{0:T(128)} multiply(%pow.46, %broadcast.4414), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %constant.1718.clone.11 = f32[]{:T(128)} constant(1) + %broadcast.4413 = f32[32]{0:T(128)} broadcast(%constant.1718.clone.11), dimensions={}, metadata={op_name="broadcast.1663"} + %constant.1730.clone.96 = f32[]{:T(128)} constant(0) + %max.96 = f32[32]{0:T(128)} broadcast(%constant.1730.clone.96), dimensions={}, metadata={op_name="jit(train_step)/jit(clip)/max" stack_frame_id=0} + %iota.345 = f32[32]{0:T(128)} iota(), iota_dimension=0, metadata={op_name="jit(train_step)/iota" stack_frame_id=0} + %constant.1739.clone.1 = f32[]{:T(128)} constant(-10) + %broadcast.4412 = f32[32]{0:T(128)} broadcast(%constant.1739.clone.1), dimensions={}, metadata={op_name="broadcast.1666"} + %sub.875 = f32[32]{0:T(128)} add(%iota.345, %broadcast.4412), metadata={op_name="jit(train_step)/sub" stack_frame_id=0} + %constant.1740.clone.1 = f32[]{:T(128)} constant(0.0769230798) + %broadcast.4411 = f32[32]{0:T(128)} broadcast(%constant.1740.clone.1), dimensions={}, metadata={op_name="broadcast.1667"} + %div.2441 = f32[32]{0:T(128)} multiply(%sub.875, %broadcast.4411), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %min.42 = f32[32]{0:T(128)} clamp(%max.96, %div.2441, %broadcast.4413), metadata={op_name="jit(train_step)/jit(clip)/min" stack_frame_id=0} + %sub.874 = f32[32]{0:T(128)} subtract(%broadcast.4413, %min.42), metadata={op_name="jit(train_step)/sub" stack_frame_id=0} + %sub.873 = f32[32]{0:T(128)} subtract(%broadcast.4413, %sub.874), metadata={op_name="jit(train_step)/sub" stack_frame_id=0} + %mul.4547 = f32[32]{0:T(128)} multiply(%div.2442, %sub.873), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.4546 = f32[32]{0:T(128)} multiply(%pow.46, %sub.874), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + ROOT %add.3632 = f32[32]{0:T(128)S(1)} add(%mul.4547, %mul.4546), metadata={op_name="jit(train_step)/add" stack_frame_id=0} +} + +%fused_computation.495 (param_0.4495: f32[32]) -> (f32[163840,32], f32[163840,32]) { + %mul.3607 = f32[163840,32]{1,0:T(8,128)} iota(), iota_dimension=0, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %param_0.4495 = f32[32]{0:T(128)S(1)} parameter(0) + %mul.3700 = f32[163840,32]{1,0:T(8,128)} broadcast(%param_0.4495), dimensions={1}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %mul.3606 = f32[163840,32]{1,0:T(8,128)} multiply(%mul.3607, %mul.3700), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %constant.1730.clone.95 = f32[]{:T(128)} constant(0) + %convert_element_type.2546 = f32[163840,32]{1,0:T(8,128)} broadcast(%constant.1730.clone.95), dimensions={}, metadata={op_name="jit(train_step)/convert_element_type" stack_frame_id=0} + %exp.480 = pred[163840,32]{1,0:T(8,128)(4,1)} compare(%mul.3606, %convert_element_type.2546), direction=EQ, metadata={op_name="jit(train_step)/exp" stack_frame_id=0} + %mul.3587 = f32[163840,32]{1,0:T(8,128)} multiply(%mul.3606, %convert_element_type.2546), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %exp.473 = f32[163840,32]{1,0:T(8,128)} exponential(%mul.3587), metadata={op_name="jit(train_step)/exp" stack_frame_id=0} + %constant.1741.clone.1 = f32[]{:T(128)} constant(inf) + %broadcast.4090 = f32[163840,32]{1,0:T(8,128)} broadcast(%constant.1741.clone.1), dimensions={}, metadata={op_name="broadcast.373"} + %exp.472 = pred[163840,32]{1,0:T(8,128)(4,1)} compare(%exp.473, %broadcast.4090), direction=EQ, metadata={op_name="jit(train_step)/exp" stack_frame_id=0} + %exp.460 = f32[163840,32]{1,0:T(8,128)} sine(%mul.3606), metadata={op_name="jit(train_step)/exp" stack_frame_id=0} + %exp.459 = f32[163840,32]{1,0:T(8,128)} multiply(%exp.473, %exp.460), metadata={op_name="jit(train_step)/exp" stack_frame_id=0} + %exp.458 = f32[163840,32]{1,0:T(8,128)} multiply(%exp.459, %exp.473), metadata={op_name="jit(train_step)/exp" stack_frame_id=0} + %exp.457 = f32[163840,32]{1,0:T(8,128)} select(%exp.472, %exp.458, %exp.459), metadata={op_name="jit(train_step)/exp" stack_frame_id=0} + %exp.453 = f32[163840,32]{1,0:T(8,128)} select(%exp.480, %convert_element_type.2546, %exp.457), metadata={op_name="jit(train_step)/exp" stack_frame_id=0} + %exp.461.clone.1 = f32[163840,32]{1,0:T(8,128)} cosine(%mul.3606), metadata={op_name="jit(train_step)/exp" stack_frame_id=0} + %exp.452.clone.1 = f32[163840,32]{1,0:T(8,128)} multiply(%exp.473, %exp.461.clone.1), metadata={op_name="jit(train_step)/exp" stack_frame_id=0} + %exp.451.clone.1 = f32[163840,32]{1,0:T(8,128)} multiply(%exp.452.clone.1, %exp.473), metadata={op_name="jit(train_step)/exp" stack_frame_id=0} + %exp.450.clone.1 = f32[163840,32]{1,0:T(8,128)} select(%exp.472, %exp.451.clone.1, %exp.452.clone.1), metadata={op_name="jit(train_step)/exp" stack_frame_id=0} + ROOT %tuple.722 = (f32[163840,32]{1,0:T(8,128)}, f32[163840,32]{1,0:T(8,128)}) tuple(%exp.453, %exp.450.clone.1) +} + +%fused_computation.1356 (param_0.3839: s32[1,128]) -> s32[128] { + %param_0.3839 = s32[1,128]{1,0:T(1,128)S(1)} parameter(0) + %constant.1719.clone.92 = s32[]{:T(128)} constant(0), metadata={op_name="jit(train_step)/shard_map/jit(gmm)" stack_frame_id=0} + %broadcast.4422 = s32[1,128]{1,0:T(1,128)} broadcast(%constant.1719.clone.92), dimensions={}, metadata={op_name="broadcast.370"} + %lt.834 = pred[1,128]{1,0:T(4,128)(4,1)} compare(%param_0.3839, %broadcast.4422), direction=LT, metadata={op_name="jit(train_step)/lt" stack_frame_id=0} + %constant.1742.clone.2 = s32[]{:T(128)} constant(163840) + %broadcast.4410 = s32[1,128]{1,0:T(1,128)} broadcast(%constant.1742.clone.2), dimensions={}, metadata={op_name="broadcast.361"} + %add.3591 = s32[1,128]{1,0:T(1,128)} add(%param_0.3839, %broadcast.4410), metadata={op_name="jit(train_step)/add" stack_frame_id=0} + %select_n.2274 = s32[1,128]{1,0:T(1,128)} select(%lt.834, %add.3591, %param_0.3839), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} + ROOT %bitcast.1929 = s32[128]{0:T(128)S(1)} bitcast(%select_n.2274), metadata={op_name="jit(train_step)/broadcast_in_dim" stack_frame_id=0} +} + +%fused_computation.1132 (param_0.3322: s32[128]) -> s32[1024] { + %constant.4112 = s32[] constant(0), metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} + %broadcast.4268 = s32[1024]{0:T(1024)} broadcast(%constant.4112), dimensions={}, metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} + %param_0.3322 = s32[128]{0:T(128)S(1)} parameter(0) + %constant.4127 = s32[] constant(2147483647), metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} + %pad.331 = s32[1024]{0:T(1024)} pad(%param_0.3322, %constant.4127), padding=0_896, metadata={op_name="jit(train_step)/gather" stack_frame_id=0} + %constant.4110 = s32[] constant(163839), metadata={op_name="jit(train_step)/gather" stack_frame_id=0} + %broadcast.4257 = s32[1024]{0:T(1024)} broadcast(%constant.4110), dimensions={}, metadata={op_name="jit(train_step)/gather" stack_frame_id=0} + ROOT %clamp.59 = s32[1024]{0:T(1024)S(1)} clamp(%broadcast.4268, %pad.331, %broadcast.4257), metadata={op_name="jit(train_step)/gather" stack_frame_id=0} +} + +%fused_computation.7 (param_0.23: f32[163840,32], param_1.123: s32[1024]) -> f32[128,32] { + %param_0.23 = f32[163840,32]{1,0:T(8,128)} parameter(0) + %param_1.123 = s32[1024]{0:T(1024)S(1)} parameter(1) + %custom-call.20 = s32[1024]{0:T(1024)} custom-call(%param_1.123), custom_call_target="AssumeGatherIndicesInBound", operand_layout_constraints={s32[1024]{0:T(1024)}}, metadata={op_name="jit(train_step)/gather" stack_frame_id=0} + %slice.1059 = s32[128]{0:T(128)} slice(%custom-call.20), slice={[0:128]}, metadata={op_name="jit(train_step)/gather" stack_frame_id=0} + %reshape.4334 = s32[128]{0:T(128)} reshape(%slice.1059), metadata={op_name="jit(train_step)/broadcast_in_dim" stack_frame_id=0} + %transpose.897 = s32[128]{0:T(128)} transpose(%reshape.4334), dimensions={0}, metadata={op_name="jit(train_step)/broadcast_in_dim" stack_frame_id=0} + %gather.228 = f32[128,32]{1,0:T(8,128)} gather(%param_0.23, %transpose.897), offset_dims={1}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1,32}, metadata={op_name="jit(train_step)/gather" stack_frame_id=0} + %transpose.896 = f32[128,32]{1,0:T(8,128)} transpose(%gather.228), dimensions={0,1}, metadata={op_name="jit(train_step)/gather" stack_frame_id=0} + ROOT %reshape.4333 = f32[128,32]{1,0:T(8,128)S(1)} reshape(%transpose.896), metadata={op_name="jit(train_step)/gather" stack_frame_id=0} } -%region_106.124 (reduce-window.204: s32[], reduce-window.205: s32[]) -> s32[] { - %reduce-window.204 = s32[]{:T(128)} parameter(0), metadata={op_name="reduce-window.71"} - %reduce-window.205 = s32[]{:T(128)} parameter(1), metadata={op_name="reduce-window.71"} - ROOT %reduce_window_sum.144 = s32[]{:T(128)} add(%reduce-window.204, %reduce-window.205), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +%fused_computation.8 (param_0.26: f32[163840,32], param_1.125: s32[1024]) -> f32[128,32] { + %param_0.26 = f32[163840,32]{1,0:T(8,128)} parameter(0) + %param_1.125 = s32[1024]{0:T(1024)S(1)} parameter(1) + %custom-call.22 = s32[1024]{0:T(1024)} custom-call(%param_1.125), custom_call_target="AssumeGatherIndicesInBound", operand_layout_constraints={s32[1024]{0:T(1024)}}, metadata={op_name="jit(train_step)/gather" stack_frame_id=0} + %slice.1061 = s32[128]{0:T(128)} slice(%custom-call.22), slice={[0:128]}, metadata={op_name="jit(train_step)/gather" stack_frame_id=0} + %reshape.4342 = s32[128]{0:T(128)} reshape(%slice.1061), metadata={op_name="jit(train_step)/broadcast_in_dim" stack_frame_id=0} + %transpose.903 = s32[128]{0:T(128)} transpose(%reshape.4342), dimensions={0}, metadata={op_name="jit(train_step)/broadcast_in_dim" stack_frame_id=0} + %gather.230 = f32[128,32]{1,0:T(8,128)} gather(%param_0.26, %transpose.903), offset_dims={1}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1,32}, metadata={op_name="jit(train_step)/gather" stack_frame_id=0} + %transpose.902 = f32[128,32]{1,0:T(8,128)} transpose(%gather.230), dimensions={0,1}, metadata={op_name="jit(train_step)/gather" stack_frame_id=0} + ROOT %reshape.4341 = f32[128,32]{1,0:T(8,128)S(1)} reshape(%transpose.902), metadata={op_name="jit(train_step)/gather" stack_frame_id=0} } -%fused_computation.5 (param_0.17: bf16[129280,512], param_1.108: s32[1024]) -> bf16[512,512] { - %param_0.17 = bf16[129280,512]{1,0:T(8,128)(2,1)} parameter(0) - %param_1.108 = s32[1024]{0:T(1024)S(1)} parameter(1) - %custom-call.13 = s32[1024]{0:T(1024)} custom-call(%param_1.108), custom_call_target="AssumeGatherIndicesInBound", operand_layout_constraints={s32[1024]{0:T(1024)}}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} - %slice.892 = s32[512]{0:T(512)} slice(%custom-call.13), slice={[0:512]}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} - %reshape.3306 = s32[4,128]{1,0:T(4,128)} reshape(%slice.892), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/broadcast_in_dim" stack_frame_id=0} - %transpose.847 = s32[4,128]{1,0:T(4,128)} transpose(%reshape.3306), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/broadcast_in_dim" stack_frame_id=0} - %gather.183 = bf16[4,128,512]{2,1,0:T(8,128)(2,1)} gather(%param_0.17, %transpose.847), offset_dims={2}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=2, slice_sizes={1,512}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} - %transpose.846 = bf16[4,128,512]{2,1,0:T(8,128)(2,1)} transpose(%gather.183), dimensions={0,1,2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} - ROOT %reshape.3305 = bf16[512,512]{1,0:T(8,128)(2,1)S(1)} reshape(%transpose.846), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} +%fused_computation.621 (param_0.1794: f32[128,3,128,128]) -> bf16[3,128,128,128] { + %param_0.1794 = f32[128,3,128,128]{3,2,1,0:T(8,128)} parameter(0) + %copy.1457 = bf16[128,3,128,128]{3,2,0,1:T(8,128)(2,1)} copy(%param_0.1794), sharding={devices=[1,1,1,4]<=[4]}, metadata={op_name="state[\'model\'][\'decoder\'][\'dense_layers\'][\'self_attention\'][\'out\'][\'kernel\'].value"} + ROOT %bitcast.1249 = bf16[3,128,128,128]{3,2,1,0:T(8,128)(2,1)} bitcast(%copy.1457), metadata={op_name="jit(train_step)/jvp()/transpose" stack_frame_id=0} } -%fused_computation.6 (param_0.20: f32[163840,32], param_1.110: s32[1024]) -> f32[512,32] { - %param_0.20 = f32[163840,32]{1,0:T(8,128)} parameter(0) - %param_1.110 = s32[1024]{0:T(1024)S(1)} parameter(1) - %custom-call.15 = s32[1024]{0:T(1024)} custom-call(%param_1.110), custom_call_target="AssumeGatherIndicesInBound", operand_layout_constraints={s32[1024]{0:T(1024)}}, metadata={op_name="jit(train_step)/dense_layers/gather" stack_frame_id=0} - %slice.894 = s32[512]{0:T(512)} slice(%custom-call.15), slice={[0:512]}, metadata={op_name="jit(train_step)/dense_layers/gather" stack_frame_id=0} - %reshape.3314 = s32[4,128]{1,0:T(4,128)} reshape(%slice.894), metadata={op_name="jit(train_step)/dense_layers/broadcast_in_dim" stack_frame_id=0} - %transpose.853 = s32[4,128]{1,0:T(4,128)} transpose(%reshape.3314), dimensions={0,1}, metadata={op_name="jit(train_step)/dense_layers/broadcast_in_dim" stack_frame_id=0} - %gather.185 = f32[4,128,32]{2,1,0:T(8,128)} gather(%param_0.20, %transpose.853), offset_dims={2}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=2, slice_sizes={1,32}, metadata={op_name="jit(train_step)/dense_layers/gather" stack_frame_id=0} - %transpose.852 = f32[4,128,32]{2,1,0:T(8,128)} transpose(%gather.185), dimensions={0,1,2}, metadata={op_name="jit(train_step)/dense_layers/gather" stack_frame_id=0} - ROOT %reshape.3313 = f32[512,32]{1,0:T(8,128)} reshape(%transpose.852), metadata={op_name="jit(train_step)/dense_layers/gather" stack_frame_id=0} +%fused_computation.1423 (param_0.4385: f32[128,3,18432]) -> bf16[3,128,18432] { + %param_0.4385 = f32[128,3,18432]{2,0,1:T(8,128)S(1)} parameter(0) + %bitcast.2154 = f32[3,128,18432]{2,1,0:T(8,128)} bitcast(%param_0.4385), metadata={op_name="jit(train_step)/jvp()/transpose" stack_frame_id=0} + ROOT %convert.634 = bf16[3,128,18432]{2,1,0:T(8,128)(2,1)} convert(%bitcast.2154) } -%fused_computation.7 (param_0.23: f32[163840,32], param_1.112: s32[1024]) -> f32[512,32] { - %param_0.23 = f32[163840,32]{1,0:T(8,128)} parameter(0) - %param_1.112 = s32[1024]{0:T(1024)S(1)} parameter(1) - %custom-call.17 = s32[1024]{0:T(1024)} custom-call(%param_1.112), custom_call_target="AssumeGatherIndicesInBound", operand_layout_constraints={s32[1024]{0:T(1024)}}, metadata={op_name="jit(train_step)/dense_layers/gather" stack_frame_id=0} - %slice.896 = s32[512]{0:T(512)} slice(%custom-call.17), slice={[0:512]}, metadata={op_name="jit(train_step)/dense_layers/gather" stack_frame_id=0} - %reshape.3322 = s32[4,128]{1,0:T(4,128)} reshape(%slice.896), metadata={op_name="jit(train_step)/dense_layers/broadcast_in_dim" stack_frame_id=0} - %transpose.859 = s32[4,128]{1,0:T(4,128)} transpose(%reshape.3322), dimensions={0,1}, metadata={op_name="jit(train_step)/dense_layers/broadcast_in_dim" stack_frame_id=0} - %gather.187 = f32[4,128,32]{2,1,0:T(8,128)} gather(%param_0.23, %transpose.859), offset_dims={2}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=2, slice_sizes={1,32}, metadata={op_name="jit(train_step)/dense_layers/gather" stack_frame_id=0} - %transpose.858 = f32[4,128,32]{2,1,0:T(8,128)} transpose(%gather.187), dimensions={0,1,2}, metadata={op_name="jit(train_step)/dense_layers/gather" stack_frame_id=0} - ROOT %reshape.3321 = f32[512,32]{1,0:T(8,128)} reshape(%transpose.858), metadata={op_name="jit(train_step)/dense_layers/gather" stack_frame_id=0} -} +%fused_computation.1421 (param_0.4384: f32[128,3,18432]) -> bf16[3,128,18432] { + %param_0.4384 = f32[128,3,18432]{2,0,1:T(8,128)S(1)} parameter(0) + %bitcast.2153 = f32[3,128,18432]{2,1,0:T(8,128)} bitcast(%param_0.4384), metadata={op_name="jit(train_step)/jvp()/transpose" stack_frame_id=0} + ROOT %convert.630 = bf16[3,128,18432]{2,1,0:T(8,128)(2,1)} convert(%bitcast.2153) +} + +%fused_computation.1422 (param_0.4383: f32[18432,3,128]) -> bf16[3,18432,128] { + %param_0.4383 = f32[18432,3,128]{2,0,1:T(8,128)} parameter(0) + %bitcast.2152 = f32[3,18432,128]{2,1,0:T(8,128)} bitcast(%param_0.4383), metadata={op_name="jit(train_step)/jvp()/transpose" stack_frame_id=0} + ROOT %convert.632 = bf16[3,18432,128]{2,1,0:T(8,128)(2,1)} convert(%bitcast.2152) +} -%fused_computation.8 (param_0.26: f32[163840,32], param_1.120: s32[1024]) -> f32[512,32] { - %param_0.26 = f32[163840,32]{1,0:T(8,128)} parameter(0) - %param_1.120 = s32[1024]{0:T(1024)S(1)} parameter(1) - %custom-call.25 = s32[1024]{0:T(1024)} custom-call(%param_1.120), custom_call_target="AssumeGatherIndicesInBound", operand_layout_constraints={s32[1024]{0:T(1024)}}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/gather" stack_frame_id=0} - %slice.904 = s32[512]{0:T(512)} slice(%custom-call.25), slice={[0:512]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/gather" stack_frame_id=0} - %reshape.3330 = s32[4,128]{1,0:T(4,128)} reshape(%slice.904), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/select_n" stack_frame_id=0} - %transpose.865 = s32[4,128]{1,0:T(4,128)} transpose(%reshape.3330), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/select_n" stack_frame_id=0} - %gather.189 = f32[4,128,32]{2,1,0:T(8,128)} gather(%param_0.26, %transpose.865), offset_dims={2}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=2, slice_sizes={1,32}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/gather" stack_frame_id=0} - %transpose.864 = f32[4,128,32]{2,1,0:T(8,128)} transpose(%gather.189), dimensions={0,1,2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/gather" stack_frame_id=0} - ROOT %reshape.3329 = f32[512,32]{1,0:T(8,128)S(1)} reshape(%transpose.864), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/gather" stack_frame_id=0} +%fused_computation.171.clone.1 (param_0.4838: bf16[3,1,128,512], param_1.5447: s32[], param_2.4298: bf16[1,128,512]) -> bf16[3,1,128,512] { + %param_0.4838 = bf16[3,1,128,512]{3,2,1,0:T(8,128)(2,1)} parameter(0) + %param_2.4298 = bf16[1,128,512]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %bitcast.2534 = bf16[1,1,128,512]{3,2,1,0:T(8,128)(2,1)} bitcast(%param_2.4298), metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_update_index_in_dim)/broadcast_in_dim" stack_frame_id=0} + %param_1.5447 = s32[]{:T(128)S(6)} parameter(1) + %constant.1492.clone.42 = s32[]{:T(128)} constant(0) + ROOT %dynamic-update-slice.65 = bf16[3,1,128,512]{3,2,1,0:T(8,128)(2,1)} dynamic-update-slice(%param_0.4838, %bitcast.2534, %param_1.5447, %constant.1492.clone.42, %constant.1492.clone.42, /*index=5*/%constant.1492.clone.42), metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_update_index_in_dim)/dynamic_update_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} } -%fused_computation.9 (param_0.29: f32[163840,32], param_1.122: s32[1024]) -> f32[512,32] { +%fused_computation.129.clone.1 (param_0.4870: bf16[3,128,128,128], param_1.5468: s32[]) -> bf16[1,128,128,128] { + %param_0.4870 = bf16[3,128,128,128]{3,2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.5468 = s32[]{:T(128)S(6)} parameter(1) + %constant.1492.clone.50 = s32[]{:T(128)} constant(0) + ROOT %dynamic-slice.532 = bf16[1,128,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} dynamic-slice(%param_0.4870, %param_1.5468, %constant.1492.clone.50, %constant.1492.clone.50, %constant.1492.clone.50), dynamic_slice_sizes={1,128,128,128}, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_index_in_dim)/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} +} + +%fused_computation.99.clone.1 (param_0.4864: bf16[3,384,128,192], param_1.5463: s32[]) -> bf16[1,384,128,192] { + %param_0.4864 = bf16[3,384,128,192]{2,1,3,0:T(8,128)(2,1)} parameter(0) + %param_1.5463 = s32[]{:T(128)S(6)} parameter(1) + %constant.1492.clone.49 = s32[]{:T(128)} constant(0) + ROOT %dynamic-slice.531 = bf16[1,384,128,192]{2,1,3,0:T(8,128)(2,1)S(1)} dynamic-slice(%param_0.4864, %param_1.5463, %constant.1492.clone.49, %constant.1492.clone.49, %constant.1492.clone.49), dynamic_slice_sizes={1,384,128,192}, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_index_in_dim)/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} +} + +%fused_computation.140.clone.1 (param_0.4859: bf16[3,128,1536], param_1.5460: s32[]) -> bf16[1,128,1536] { + %param_0.4859 = bf16[3,128,1536]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.5460 = s32[]{:T(128)S(6)} parameter(1) + %constant.1492.clone.48 = s32[]{:T(128)} constant(0) + ROOT %dynamic-slice.530 = bf16[1,128,1536]{2,1,0:T(8,128)(2,1)S(1)} dynamic-slice(%param_0.4859, %param_1.5460, %constant.1492.clone.48, %constant.1492.clone.48), dynamic_slice_sizes={1,128,1536}, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_index_in_dim)/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} +} + +%region_2.3 (reduce_sum.9: f32[], reduce_sum.13: f32[]) -> f32[] { + %reduce_sum.9 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/reduce_sum"} + %reduce_sum.13 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/reduce_sum"} + ROOT %reduce_sum.14 = f32[]{:T(128)} add(%reduce_sum.9, %reduce_sum.13), metadata={op_name="reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%fused_computation.158.clone.1 (param_0.4841: bf16[1,128,512]) -> f32[128] { + %param_0.4841 = bf16[1,128,512]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.3009 = f32[1,128,512]{2,1,0:T(8,128)} convert(%param_0.4841), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %square.699 = f32[1,128,512]{2,1,0:T(8,128)} multiply(%convert_element_type.3009, %convert_element_type.3009), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/square" stack_frame_id=0} + %constant.1493.clone.19 = f32[]{:T(128)} constant(0) + ROOT %reduce.880 = f32[128]{0:T(128)S(1)} reduce(%square.699, %constant.1493.clone.19), dimensions={0,2}, to_apply=%region_2.3, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/reduce_sum" stack_frame_id=0} +} + +%fused_computation.205.clone.1 (param_0.4842: f32[128]) -> f32[128] { + %param_0.4842 = f32[128]{0:T(128)S(1)} parameter(0) + %constant.1494.clone.11 = f32[]{:T(128)} constant(0.001953125) + %broadcast.4707 = f32[128]{0:T(128)} broadcast(%constant.1494.clone.11), dimensions={}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call" stack_frame_id=0} + %div.2870 = f32[128]{0:T(128)} multiply(%param_0.4842, %broadcast.4707), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/div" stack_frame_id=0} + %constant.1495.clone.13 = f32[]{:T(128)} constant(1e-06) + %broadcast.4706 = f32[128]{0:T(128)} broadcast(%constant.1495.clone.13), dimensions={}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call" stack_frame_id=0} + %add.3967 = f32[128]{0:T(128)} add(%div.2870, %broadcast.4706), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/add" stack_frame_id=0} + ROOT %rsqrt.366 = f32[128]{0:T(128)S(1)} rsqrt(%add.3967), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/rsqrt" stack_frame_id=0} +} + +%convert_element_type.2195.reduce_sub_computation (lhs.15: bf16[], rhs.15: bf16[]) -> bf16[] { + %lhs.15 = bf16[] parameter(0) + %rhs.15 = bf16[] parameter(1) + ROOT %add.3047 = bf16[] add(%lhs.15, %rhs.15), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%convert_element_type.2198.reduce_sub_computation (lhs.17: bf16[], rhs.17: bf16[]) -> bf16[] { + %lhs.17 = bf16[] parameter(0) + %rhs.17 = bf16[] parameter(1) + ROOT %add.3049 = bf16[] add(%lhs.17, %rhs.17), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%fused_computation.197.clone.1 (param_0.4839: bf16[3,512], param_1.5448: s32[], param_2.4299: bf16[3,512]) -> (bf16[512], bf16[512]) { + %param_0.4839 = bf16[3,512]{1,0:T(4,128)(2,1)} parameter(0) + %param_1.5448 = s32[]{:T(128)S(6)} parameter(1) + %constant.1492.clone.43 = s32[]{:T(128)} constant(0) + %dynamic_slice.628 = bf16[1,512]{1,0:T(2,128)(2,1)} dynamic-slice(%param_0.4839, %param_1.5448, %constant.1492.clone.43), dynamic_slice_sizes={1,512}, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_index_in_dim)/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %constant.4256 = bf16[]{:T(256)} constant(-0), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %reduce.879 = bf16[512]{0:T(512)(128)(2,1)S(1)} reduce(%dynamic_slice.628, %constant.4256), dimensions={0}, to_apply=%convert_element_type.2195.reduce_sub_computation, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %param_2.4299 = bf16[3,512]{1,0:T(4,128)(2,1)} parameter(2) + %dynamic_slice.614.clone.3 = bf16[1,512]{1,0:T(2,128)(2,1)} dynamic-slice(%param_2.4299, %param_1.5448, %constant.1492.clone.43), dynamic_slice_sizes={1,512}, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_index_in_dim)/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %reduce.682.clone.3 = bf16[512]{0:T(512)(128)(2,1)S(1)} reduce(%dynamic_slice.614.clone.3, %constant.4256), dimensions={0}, to_apply=%convert_element_type.2198.reduce_sub_computation, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + ROOT %tuple.862 = (bf16[512]{0:T(512)(128)(2,1)S(1)}, bf16[512]{0:T(512)(128)(2,1)S(1)}) tuple(%reduce.879, %reduce.682.clone.3) +} + +%fused_computation.170.clone.1.clone.clone.clone.1 (param_0.4861: bf16[1,128,512], param_1.5461: f32[128], param_2.4305: bf16[512]) -> bf16[128,512] { + %param_2.4305 = bf16[512]{0:T(512)(128)(2,1)S(1)} parameter(2) + %dot_general.918 = bf16[1,128,512]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.4305), dimensions={2}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/...k,k->...k/dot_general" stack_frame_id=0} + %convert.1035 = f32[1,128,512]{2,1,0:T(8,128)} convert(%dot_general.918) + %param_0.4861 = bf16[1,128,512]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.3019 = f32[1,128,512]{2,1,0:T(8,128)} convert(%param_0.4861), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %param_1.5461 = f32[128]{0:T(128)S(1)} parameter(1) + %mul.5148 = f32[1,128,512]{2,1,0:T(8,128)} broadcast(%param_1.5461), dimensions={1}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %mul.5147 = f32[1,128,512]{2,1,0:T(8,128)} multiply(%convert_element_type.3019, %mul.5148), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %convert_element_type.3018 = bf16[1,128,512]{2,1,0:T(8,128)(2,1)} convert(%mul.5147), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %convert.1036 = f32[1,128,512]{2,1,0:T(8,128)} convert(%convert_element_type.3018) + %dot_general.917 = f32[1,128,512]{2,1,0:T(8,128)} multiply(%convert.1035, %convert.1036), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/...k,k->...k/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","512"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %convert.1037 = bf16[1,128,512]{2,1,0:T(8,128)(2,1)} convert(%dot_general.917) + ROOT %bitcast.2549 = bf16[128,512]{1,0:T(8,128)(2,1)} bitcast(%convert.1037) +} + +%fused_computation.135.clone.clone.clone.clone.clone.clone.clone.1 (param_0.4860: bf16[1,512,1536]) -> bf16[512,1536] { + %param_0.4860 = bf16[1,512,1536]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.2548 = bf16[512,1536]{1,0:T(8,128)(2,1)} bitcast(%param_0.4860), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} +} + +%region_3.4 (reduce_sum.15: f32[], reduce_sum.16: f32[]) -> f32[] { + %reduce_sum.15 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/reduce_sum"} + %reduce_sum.16 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/reduce_sum"} + ROOT %reduce_sum.20 = f32[]{:T(128)} add(%reduce_sum.15, %reduce_sum.16), metadata={op_name="reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%fused_computation.139.clone.1 (param_0.4862: bf16[1,512,1536], param_1.5462: bf16[1,128,512], param_2.4306: f32[128], param_3.2900: bf16[512]) -> (f32[128], f32[1,128,1536]) { + %param_1.5462 = bf16[1,128,512]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %param_2.4306 = f32[128]{0:T(128)S(1)} parameter(2) + %param_3.2900 = bf16[512]{0:T(512)(128)(2,1)S(1)} parameter(3) + %fusion.140.clone.3 = bf16[128,512]{1,0:T(8,128)(2,1)} fusion(%param_1.5462, %param_2.4306, %param_3.2900), kind=kLoop, calls=%fused_computation.170.clone.1.clone.clone.clone.1 + %param_0.4862 = bf16[1,512,1536]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %fusion.85.clone.3 = bf16[512,1536]{1,0:T(8,128)(2,1)} fusion(%param_0.4862), kind=kLoop, calls=%fused_computation.135.clone.clone.clone.clone.clone.clone.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} + %convolution.91.clone.3 = bf16[128,1536]{1,0:T(8,128)(2,1)} convolution(%fusion.140.clone.3, %fusion.85.clone.3), dim_labels=bf_io->bf, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} + %convert_element_type.2255.clone.3 = f32[128,1536]{1,0:T(8,128)} convert(%convolution.91.clone.3), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %bitcast.695.clone.3 = f32[1,128,1536]{2,1,0:T(8,128)S(1)} bitcast(%convert_element_type.2255.clone.3), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %square.701 = f32[1,128,1536]{2,1,0:T(8,128)} multiply(%bitcast.695.clone.3, %bitcast.695.clone.3), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/square" stack_frame_id=0} + %constant.1493.clone.22 = f32[]{:T(128)} constant(0) + %reduce.882 = f32[128]{0:T(128)S(1)} reduce(%square.701, %constant.1493.clone.22), dimensions={0,2}, to_apply=%region_3.4, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/reduce_sum" stack_frame_id=0} + ROOT %tuple.866 = (f32[128]{0:T(128)S(1)}, f32[1,128,1536]{2,1,0:T(8,128)S(1)}) tuple(%reduce.882, %bitcast.695.clone.3) +} + +%fused_computation.204.clone.1 (param_0.4863: f32[128]) -> f32[128] { + %param_0.4863 = f32[128]{0:T(128)S(1)} parameter(0) + %constant.1496.clone.7 = f32[]{:T(128)} constant(0.000651041686) + %broadcast.4711 = f32[128]{0:T(128)} broadcast(%constant.1496.clone.7), dimensions={}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/div" stack_frame_id=0} + %div.2872 = f32[128]{0:T(128)} multiply(%param_0.4863, %broadcast.4711), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/div" stack_frame_id=0} + %constant.1495.clone.15 = f32[]{:T(128)} constant(1e-06) + %broadcast.4710 = f32[128]{0:T(128)} broadcast(%constant.1495.clone.15), dimensions={}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call" stack_frame_id=0} + %add.3970 = f32[128]{0:T(128)} add(%div.2872, %broadcast.4710), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/add" stack_frame_id=0} + ROOT %rsqrt.368 = f32[128]{0:T(128)S(1)} rsqrt(%add.3970), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/rsqrt" stack_frame_id=0} +} + +%fused_computation.200.clone.1 (param_0.4858: bf16[3,384], param_1.5459: s32[]) -> bf16[1,384] { + %param_0.4858 = bf16[3,384]{1,0:T(4,128)(2,1)} parameter(0) + %param_1.5459 = s32[]{:T(128)S(6)} parameter(1) + %constant.1492.clone.47 = s32[]{:T(128)} constant(0) + ROOT %dynamic-slice.529 = bf16[1,384]{1,0:T(2,128)(2,1)S(1)} dynamic-slice(%param_0.4858, %param_1.5459, %constant.1492.clone.47), dynamic_slice_sizes={1,384}, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_index_in_dim)/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} +} + +%all-gather.290.reduce_sub_computation (lhs.16: bf16[], rhs.16: bf16[]) -> bf16[] { + %lhs.16 = bf16[] parameter(0) + %rhs.16 = bf16[] parameter(1) + ROOT %add.3048 = bf16[] add(%lhs.16, %rhs.16), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%fused_computation.149.clone.clone.1 (param_0.4865: f32[1,128,1536], param_1.5464: f32[128], param_2.4307: bf16[1536]) -> bf16[128,1536,1] { + %param_2.4307 = bf16[1536]{0:T(1024)(128)(2,1)S(1)} parameter(2) + %dot_general.920 = bf16[1,128,1536]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.4307), dimensions={2}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/...k,k->...k/dot_general" stack_frame_id=0} + %convert.1038 = f32[1,128,1536]{2,1,0:T(8,128)} convert(%dot_general.920) + %param_0.4865 = f32[1,128,1536]{2,1,0:T(8,128)S(1)} parameter(0) + %param_1.5464 = f32[128]{0:T(128)S(1)} parameter(1) + %mul.5150 = f32[1,128,1536]{2,1,0:T(8,128)} broadcast(%param_1.5464), dimensions={1}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %mul.5149 = f32[1,128,1536]{2,1,0:T(8,128)} multiply(%param_0.4865, %mul.5150), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %convert_element_type.3020 = bf16[1,128,1536]{2,1,0:T(8,128)(2,1)} convert(%mul.5149), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %convert.1039 = f32[1,128,1536]{2,1,0:T(8,128)} convert(%convert_element_type.3020) + %dot_general.919 = f32[1,128,1536]{2,1,0:T(8,128)} multiply(%convert.1038, %convert.1039), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/...k,k->...k/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","1536"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %convert.1040 = bf16[1,128,1536]{2,1,0:T(8,128)(2,1)} convert(%dot_general.919) + ROOT %bitcast.2551 = bf16[128,1536,1]{1,0,2:T(8,128)(2,1)} bitcast(%convert.1040) +} + +%bitcast_fusion.17.clone.1 (bitcast_input.27: bf16[1536,128,192]) -> bf16[1536,128,192] { + %bitcast_input.27 = bf16[1536,128,192]{0,2,1:T(8,128)(2,1)} parameter(0) + ROOT %bitcast.2550 = bf16[1536,128,192]{0,2,1:T(8,128)(2,1)} bitcast(%bitcast_input.27) +} + +%fused_computation.131.clone.1 (param_0.4866: bf16[1536,128,192], param_1.5465: f32[1,128,1536], param_2.4308: f32[128], param_3.2901: bf16[1536]) -> bf16[1,128,128,192] { + %param_1.5465 = f32[1,128,1536]{2,1,0:T(8,128)S(1)} parameter(1) + %param_2.4308 = f32[128]{0:T(128)S(1)} parameter(2) + %param_3.2901 = bf16[1536]{0:T(1024)(128)(2,1)S(1)} parameter(3) + %fusion.1120 = bf16[128,1536,1]{1,0,2:T(8,128)(2,1)} fusion(%param_1.5465, %param_2.4308, %param_3.2901), kind=kLoop, calls=%fused_computation.149.clone.clone.1 + %param_0.4866 = bf16[1536,128,192]{0,2,1:T(8,128)(2,1)} parameter(0) + %fusion.1121 = bf16[1536,128,192]{0,2,1:T(8,128)(2,1)} fusion(%param_0.4866), kind=kLoop, calls=%bitcast_fusion.17.clone.1 + %convolution.334 = bf16[128,128,192]{2,0,1:T(8,128)(2,1)} convolution(%fusion.1120, %fusion.1121), window={size=128 pad=127_127 rhs_reversal=1}, dim_labels=bf0_i0o->b0f, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} + ROOT %bitcast.2552 = bf16[1,128,128,192]{3,1,2,0:T(8,128)(2,1)S(1)} bitcast(%convolution.334), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} +} + +%fused_computation.1415.clone.1 (param_0.4867: bf16[1,128,128,32,2]) -> (bf16[1,128,128,32,1], bf16[1,128,128,32,1]) { + %param_0.4867 = bf16[1,128,128,32,2]{2,1,4,3,0:T(8,128)(2,1)S(1)} parameter(0) + %slice.1442 = bf16[1,128,128,32,1]{2,1,4,3,0:T(8,128)(2,1)S(1)} slice(%param_0.4867), slice={[0:1], [0:128], [0:128], [0:32], [0:1]}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/slice" stack_frame_id=0} + %slice.1443 = bf16[1,128,128,32,1]{2,1,4,3,0:T(8,128)(2,1)S(1)} slice(%param_0.4867), slice={[0:1], [0:128], [0:128], [0:32], [1:2]}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/slice" stack_frame_id=0} + ROOT %tuple.867 = (bf16[1,128,128,32,1]{2,1,4,3,0:T(8,128)(2,1)S(1)}, bf16[1,128,128,32,1]{2,1,4,3,0:T(8,128)(2,1)S(1)}) tuple(%slice.1442, %slice.1443) +} + +%fused_computation.126.clone.1 (param_0.4868: bf16[1,128,128,32,1], param_1.5466: bf16[1,128,128,32,1], param_2.4309: f32[128,32], param_3.2902: f32[128,32]) -> (bf16[1,128,128,32], bf16[1,128,128,32]) { + %param_1.5466 = bf16[1,128,128,32,1]{3,1,2,4,0:T(8,128)(2,1)S(1)} parameter(1) + %convert_element_type.3022 = f32[1,128,128,32,1]{3,1,2,4,0:T(8,128)} convert(%param_1.5466), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %bitcast.2554 = f32[1,128,128,32]{3,1,2,0:T(8,128)} bitcast(%convert_element_type.3022), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %param_0.4868 = bf16[1,128,128,32,1]{3,1,2,4,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.3021 = f32[1,128,128,32,1]{3,1,2,4,0:T(8,128)} convert(%param_0.4868), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %bitcast.2553 = f32[1,128,128,32]{3,1,2,0:T(8,128)} bitcast(%convert_element_type.3021), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %constant.1493.clone.23 = f32[]{:T(128)} constant(0) + %convert_element_type.3023 = f32[1,128,128,32]{3,1,2,0:T(8,128)} broadcast(%constant.1493.clone.23), dimensions={}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %mul.5154 = f32[1,128,128,32]{3,1,2,0:T(8,128)} multiply(%bitcast.2553, %convert_element_type.3023), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %add.3971 = f32[1,128,128,32]{3,1,2,0:T(8,128)} add(%bitcast.2554, %mul.5154), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/add" stack_frame_id=0} + %param_3.2902 = f32[128,32]{1,0:T(8,128)S(1)} parameter(3) + %broadcast_in_dim.1648 = f32[1,128,128,32]{3,1,2,0:T(8,128)} broadcast(%param_3.2902), dimensions={1,3}, metadata={op_name="jit(train_step)/broadcast_in_dim" stack_frame_id=0} + %mul.5153 = f32[1,128,128,32]{3,1,2,0:T(8,128)} multiply(%add.3971, %broadcast_in_dim.1648), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %param_2.4309 = f32[128,32]{1,0:T(8,128)S(1)} parameter(2) + %broadcast_in_dim.1647 = f32[1,128,128,32]{3,1,2,0:T(8,128)} broadcast(%param_2.4309), dimensions={1,3}, metadata={op_name="jit(train_step)/broadcast_in_dim" stack_frame_id=0} + %mul.5152 = f32[1,128,128,32]{3,1,2,0:T(8,128)} multiply(%bitcast.2553, %broadcast_in_dim.1647), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %mul.5151 = f32[1,128,128,32]{3,1,2,0:T(8,128)} subtract(%mul.5153, %mul.5152), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %convert.675 = bf16[1,128,128,32]{3,1,2,0:T(8,128)(2,1)S(1)} convert(%mul.5151), metadata={op_name="convert.317"} + %mul.3131.clone.3 = f32[1,128,128,32]{3,1,2,0:T(8,128)} multiply(%add.3971, %broadcast_in_dim.1647), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %mul.3130.clone.3 = f32[1,128,128,32]{3,1,2,0:T(8,128)} multiply(%bitcast.2553, %broadcast_in_dim.1648), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %mul.3129.clone.3 = f32[1,128,128,32]{3,1,2,0:T(8,128)} add(%mul.3131.clone.3, %mul.3130.clone.3), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %convert.518.clone.3 = bf16[1,128,128,32]{3,1,2,0:T(8,128)(2,1)S(1)} convert(%mul.3129.clone.3), metadata={op_name="convert.318"} + ROOT %tuple.868 = (bf16[1,128,128,32]{3,1,2,0:T(8,128)(2,1)S(1)}, bf16[1,128,128,32]{3,1,2,0:T(8,128)(2,1)S(1)}) tuple(%convert.675, %convert.518.clone.3) +} + +%fused_computation.119.clone.1 (param_0.4869: bf16[1,128,128,32], param_1.5467: bf16[1,128,128,32], param_2.4310: bf16[1,128,128,128]) -> bf16[128,128,192] { + %param_2.4310 = bf16[1,128,128,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(2) + %constant.4259 = bf16[]{:T(256)} constant(-inf) + %pad.463 = bf16[1,128,128,192]{3,1,2,0:T(8,128)(2,1)} pad(%param_2.4310, %constant.4259), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/concatenate" stack_frame_id=0} + %convert.1051 = f32[1,128,128,192]{3,1,2,0:T(8,128)} convert(%pad.463) + %param_1.5467 = bf16[1,128,128,32]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) + %pad.462 = bf16[1,128,128,192]{3,1,2,0:T(8,128)(2,1)} pad(%param_1.5467, %constant.4259), padding=0_0x0_0x0_0x128_32, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/concatenate" stack_frame_id=0} + %convert.1052 = f32[1,128,128,192]{3,1,2,0:T(8,128)} convert(%pad.462) + %maximum.100 = f32[1,128,128,192]{3,1,2,0:T(8,128)} maximum(%convert.1051, %convert.1052), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/concatenate" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","128","192"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %param_0.4869 = bf16[1,128,128,32]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %pad.461 = bf16[1,128,128,192]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.4869, %constant.4259), padding=0_0x0_0x0_0x160_0, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/concatenate" stack_frame_id=0} + %convert.1053 = f32[1,128,128,192]{3,1,2,0:T(8,128)} convert(%pad.461) + %maximum.99 = f32[1,128,128,192]{3,1,2,0:T(8,128)} maximum(%maximum.100, %convert.1053), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/concatenate" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","128","192"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %constant.1497.clone.7 = bf16[]{:T(256)} constant(0.1348) + %mul.5156 = bf16[1,128,128,192]{3,1,2,0:T(8,128)(2,1)} broadcast(%constant.1497.clone.7), dimensions={}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %convert.1054 = f32[1,128,128,192]{3,1,2,0:T(8,128)} convert(%mul.5156) + %mul.5155 = f32[1,128,128,192]{3,1,2,0:T(8,128)} multiply(%maximum.99, %convert.1054), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","128","192"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %convert.1055 = bf16[1,128,128,192]{3,1,2,0:T(8,128)(2,1)} convert(%mul.5155) + ROOT %bitcast.2555 = bf16[128,128,192]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%convert.1055), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/vmap(static_sparsity_splash)/jit(_splash_attention)/splash_mha_fwd_segmented_residuals/splash_mha_fwd_segmented_residuals/squeeze" stack_frame_id=0} +} + +%fused_computation.155.clone.1 (param_0.4843: bf16[3,128,576], param_1.5450: s32[]) -> bf16[1,128,576] { + %param_0.4843 = bf16[3,128,576]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %param_1.5450 = s32[]{:T(128)S(6)} parameter(1) + %constant.1492.clone.45 = s32[]{:T(128)} constant(0) + ROOT %dynamic-slice.527 = bf16[1,128,576]{2,1,0:T(8,128)(2,1)S(1)} dynamic-slice(%param_0.4843, %param_1.5450, %constant.1492.clone.45, %constant.1492.clone.45), dynamic_slice_sizes={1,128,576}, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_index_in_dim)/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} +} + +%fused_computation.170.clone.clone.1 (param_0.4845: bf16[1,128,512], param_1.5451: f32[128], param_2.4300: bf16[512]) -> bf16[128,512] { + %param_2.4300 = bf16[512]{0:T(512)(128)(2,1)S(1)} parameter(2) + %dot_general.914 = bf16[1,128,512]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.4300), dimensions={2}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/...k,k->...k/dot_general" stack_frame_id=0} + %convert.1029 = f32[1,128,512]{2,1,0:T(8,128)} convert(%dot_general.914) + %param_0.4845 = bf16[1,128,512]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.3011 = f32[1,128,512]{2,1,0:T(8,128)} convert(%param_0.4845), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %param_1.5451 = f32[128]{0:T(128)S(1)} parameter(1) + %mul.5140 = f32[1,128,512]{2,1,0:T(8,128)} broadcast(%param_1.5451), dimensions={1}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %mul.5139 = f32[1,128,512]{2,1,0:T(8,128)} multiply(%convert_element_type.3011, %mul.5140), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %convert_element_type.3010 = bf16[1,128,512]{2,1,0:T(8,128)(2,1)} convert(%mul.5139), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %convert.1030 = f32[1,128,512]{2,1,0:T(8,128)} convert(%convert_element_type.3010) + %dot_general.913 = f32[1,128,512]{2,1,0:T(8,128)} multiply(%convert.1029, %convert.1030), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/...k,k->...k/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","512"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %convert.1031 = bf16[1,128,512]{2,1,0:T(8,128)(2,1)} convert(%dot_general.913) + ROOT %bitcast.2536 = bf16[128,512]{1,0:T(8,128)(2,1)} bitcast(%convert.1031) +} + +%fused_computation.142.clone.clone.clone.1 (param_0.4844: bf16[1,512,576]) -> bf16[512,576] { + %param_0.4844 = bf16[1,512,576]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.2535 = bf16[512,576]{1,0:T(8,128)(2,1)} bitcast(%param_0.4844), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} +} + +%fused_computation.143.clone.1 (param_0.4846: bf16[1,512,576], param_1.5452: bf16[1,128,512], param_2.4301: f32[128], param_3.2897: bf16[512]) -> bf16[1,128,576] { + %param_1.5452 = bf16[1,128,512]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %param_2.4301 = f32[128]{0:T(128)S(1)} parameter(2) + %param_3.2897 = bf16[512]{0:T(512)(128)(2,1)S(1)} parameter(3) + %fusion.1117 = bf16[128,512]{1,0:T(8,128)(2,1)} fusion(%param_1.5452, %param_2.4301, %param_3.2897), kind=kLoop, calls=%fused_computation.170.clone.clone.1 + %param_0.4846 = bf16[1,512,576]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %fusion.1116 = bf16[512,576]{1,0:T(8,128)(2,1)} fusion(%param_0.4846), kind=kLoop, calls=%fused_computation.142.clone.clone.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} + %convolution.332 = bf16[128,576]{0,1:T(8,128)(2,1)} convolution(%fusion.1117, %fusion.1116), dim_labels=bf_io->bf, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} + ROOT %bitcast.2537 = bf16[1,128,576]{1,2,0:T(8,128)(2,1)S(1)} bitcast(%convolution.332), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} +} + +%fused_computation.1416.clone.1 (param_0.4854: bf16[1,128,1,32,2]) -> (bf16[1,128,1,32,1], bf16[1,128,1,32,1]) { + %param_0.4854 = bf16[1,128,1,32,2]{1,4,3,2,0:T(2,128)(2,1)S(1)} parameter(0) + %slice.1440 = bf16[1,128,1,32,1]{1,4,3,2,0:T(2,128)(2,1)S(1)} slice(%param_0.4854), slice={[0:1], [0:128], [0:1], [0:32], [0:1]}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/slice" stack_frame_id=0} + %slice.1441 = bf16[1,128,1,32,1]{1,4,3,2,0:T(2,128)(2,1)S(1)} slice(%param_0.4854), slice={[0:1], [0:128], [0:1], [0:32], [1:2]}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/slice" stack_frame_id=0} + ROOT %tuple.864 = (bf16[1,128,1,32,1]{1,4,3,2,0:T(2,128)(2,1)S(1)}, bf16[1,128,1,32,1]{1,4,3,2,0:T(2,128)(2,1)S(1)}) tuple(%slice.1440, %slice.1441) +} + +%fused_computation.179.clone.1 (param_0.4855: f32[1,128,1,32], param_1.5456: f32[1,128,1,32], param_2.4304: bf16[1,128,1,32,1], param_3.2899: bf16[1,128,1,32,1]) -> (bf16[1,128,32], bf16[1,128,32]) { + %param_3.2899 = bf16[1,128,1,32,1]{3,1,4,2,0:T(8,128)(2,1)S(1)} parameter(3) + %convert_element_type.3016 = f32[1,128,1,32,1]{3,1,4,2,0:T(8,128)} convert(%param_3.2899), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %bitcast.2545 = f32[1,128,32]{2,1,0:T(8,128)} bitcast(%convert_element_type.3016), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %param_2.4304 = bf16[1,128,1,32,1]{3,1,4,2,0:T(8,128)(2,1)S(1)} parameter(2) + %convert_element_type.3015 = f32[1,128,1,32,1]{3,1,4,2,0:T(8,128)} convert(%param_2.4304), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %bitcast.2544 = f32[1,128,32]{2,1,0:T(8,128)} bitcast(%convert_element_type.3015), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %constant.1493.clone.21 = f32[]{:T(128)} constant(0) + %convert_element_type.3017 = f32[1,128,32]{2,1,0:T(8,128)} broadcast(%constant.1493.clone.21), dimensions={}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %mul.5146 = f32[1,128,32]{2,1,0:T(8,128)} multiply(%bitcast.2544, %convert_element_type.3017), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %add.3969 = f32[1,128,32]{2,1,0:T(8,128)} add(%bitcast.2545, %mul.5146), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/add" stack_frame_id=0} + %param_1.5456 = f32[1,128,1,32]{3,1,2,0:T(8,128)S(1)} parameter(1) + %bitcast.2543 = f32[1,128,32]{2,1,0:T(8,128)} bitcast(%param_1.5456), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %mul.5145 = f32[1,128,32]{2,1,0:T(8,128)} multiply(%add.3969, %bitcast.2543), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %param_0.4855 = f32[1,128,1,32]{3,1,2,0:T(8,128)S(1)} parameter(0) + %bitcast.2542 = f32[1,128,32]{2,1,0:T(8,128)} bitcast(%param_0.4855), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %mul.5144 = f32[1,128,32]{2,1,0:T(8,128)} multiply(%bitcast.2544, %bitcast.2542), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %mul.5143 = f32[1,128,32]{2,1,0:T(8,128)} subtract(%mul.5145, %mul.5144), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %convert.674 = bf16[1,128,32]{2,1,0:T(8,128)(2,1)S(1)} convert(%mul.5143), metadata={op_name="convert.315"} + %mul.3219.clone.3 = f32[1,128,32]{2,1,0:T(8,128)} multiply(%add.3969, %bitcast.2542), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %mul.3225.clone.3 = f32[1,128,32]{2,1,0:T(8,128)} multiply(%bitcast.2544, %bitcast.2543), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %mul.3218.clone.3 = f32[1,128,32]{2,1,0:T(8,128)} add(%mul.3219.clone.3, %mul.3225.clone.3), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %convert.520.clone.3 = bf16[1,128,32]{2,1,0:T(8,128)(2,1)S(1)} convert(%mul.3218.clone.3), metadata={op_name="convert.316"} + ROOT %tuple.865 = (bf16[1,128,32]{2,1,0:T(8,128)(2,1)S(1)}, bf16[1,128,32]{2,1,0:T(8,128)(2,1)S(1)}) tuple(%convert.674, %convert.520.clone.3) +} + +%fused_computation.191.clone.1 (param_0.4856: bf16[1,128,32], param_1.5457: bf16[1,128,32]) -> bf16[128,64] { + %param_1.5457 = bf16[1,128,32]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %constant.4257 = bf16[]{:T(256)} constant(-inf) + %pad.458 = bf16[1,128,64]{2,1,0:T(8,128)(2,1)} pad(%param_1.5457, %constant.4257), padding=0_0x0_0x0_32, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/concatenate" stack_frame_id=0} + %convert.1056 = f32[1,128,64]{2,1,0:T(8,128)} convert(%pad.458) + %param_0.4856 = bf16[1,128,32]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %pad.457 = bf16[1,128,64]{2,1,0:T(8,128)(2,1)} pad(%param_0.4856, %constant.4257), padding=0_0x0_0x32_0, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/concatenate" stack_frame_id=0} + %convert.1057 = f32[1,128,64]{2,1,0:T(8,128)} convert(%pad.457) + %maximum.97 = f32[1,128,64]{2,1,0:T(8,128)} maximum(%convert.1056, %convert.1057), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/concatenate" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","64"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %convert.1058 = bf16[1,128,64]{2,1,0:T(8,128)(2,1)} convert(%maximum.97) + ROOT %bitcast.2546 = bf16[128,64]{1,0:T(8,128)(2,1)S(1)} bitcast(%convert.1058), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/concatenate" stack_frame_id=0} +} + +%fused_computation.110.clone.1 (param_0.4849: bf16[3,128,128,256], param_1.5453: s32[]) -> bf16[1,128,128,256] { + %param_0.4849 = bf16[3,128,128,256]{3,1,2,0:T(8,128)(2,1)} parameter(0) + %param_1.5453 = s32[]{:T(128)S(6)} parameter(1) + %constant.1492.clone.46 = s32[]{:T(128)} constant(0) + ROOT %dynamic-slice.528 = bf16[1,128,128,256]{3,1,2,0:T(8,128)(2,1)S(1)} dynamic-slice(%param_0.4849, %param_1.5453, %constant.1492.clone.46, %constant.1492.clone.46, %constant.1492.clone.46), dynamic_slice_sizes={1,128,128,256}, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_index_in_dim)/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} +} + +%region_4.5 (reduce_sum.21: f32[], reduce_sum.22: f32[]) -> f32[] { + %reduce_sum.21 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/reduce_sum"} + %reduce_sum.22 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/reduce_sum"} + ROOT %reduce_sum.23 = f32[]{:T(128)} add(%reduce_sum.21, %reduce_sum.22), metadata={op_name="reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%fused_computation.160.clone.1 (param_0.4847: bf16[1,128,576]) -> f32[128] { + %param_0.4847 = bf16[1,128,576]{1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %slice.1436 = bf16[1,128,512]{1,2,0:T(8,128)(2,1)} slice(%param_0.4847), slice={[0:1], [0:128], [0:512]}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/split" stack_frame_id=0} + %convert_element_type.3012 = f32[1,128,512]{1,2,0:T(8,128)} convert(%slice.1436), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %square.700 = f32[1,128,512]{1,2,0:T(8,128)} multiply(%convert_element_type.3012, %convert_element_type.3012), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/square" stack_frame_id=0} + %constant.1493.clone.20 = f32[]{:T(128)} constant(0) + ROOT %reduce.881 = f32[128]{0:T(128)S(1)} reduce(%square.700, %constant.1493.clone.20), dimensions={0,2}, to_apply=%region_4.5, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/reduce_sum" stack_frame_id=0} +} + +%fused_computation.203.clone.1 (param_0.4848: f32[128]) -> f32[128] { + %param_0.4848 = f32[128]{0:T(128)S(1)} parameter(0) + %constant.1494.clone.12 = f32[]{:T(128)} constant(0.001953125) + %broadcast.4709 = f32[128]{0:T(128)} broadcast(%constant.1494.clone.12), dimensions={}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call" stack_frame_id=0} + %div.2871 = f32[128]{0:T(128)} multiply(%param_0.4848, %broadcast.4709), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/div" stack_frame_id=0} + %constant.1495.clone.14 = f32[]{:T(128)} constant(1e-06) + %broadcast.4708 = f32[128]{0:T(128)} broadcast(%constant.1495.clone.14), dimensions={}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call" stack_frame_id=0} + %add.3968 = f32[128]{0:T(128)} add(%div.2871, %broadcast.4708), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/add" stack_frame_id=0} + ROOT %rsqrt.367 = f32[128]{0:T(128)S(1)} rsqrt(%add.3968), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/rsqrt" stack_frame_id=0} +} + +%fused_computation.201.clone.1 (param_0.4840: bf16[3,128], param_1.5449: s32[]) -> bf16[1,128] { + %param_0.4840 = bf16[3,128]{1,0:T(4,128)(2,1)S(1)} parameter(0) + %param_1.5449 = s32[]{:T(128)S(6)} parameter(1) + %constant.1492.clone.44 = s32[]{:T(128)} constant(0) + ROOT %dynamic-slice.526 = bf16[1,128]{1,0:T(2,128)(2,1)S(1)} dynamic-slice(%param_0.4840, %param_1.5449, %constant.1492.clone.44), dynamic_slice_sizes={1,128}, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_index_in_dim)/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} +} + +%all-gather.293.reduce_sub_computation (lhs.18: bf16[], rhs.18: bf16[]) -> bf16[] { + %lhs.18 = bf16[] parameter(0) + %rhs.18 = bf16[] parameter(1) + ROOT %add.3050 = bf16[] add(%lhs.18, %rhs.18), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%fused_computation.168.clone.clone.1 (param_0.4851: f32[128], param_1.5454: bf16[1,128,576], param_2.4302: bf16[512]) -> bf16[128,512,1] { + %param_2.4302 = bf16[512]{0:T(512)(128)(2,1)S(1)} parameter(2) + %dot_general.916 = bf16[1,128,512]{1,2,0:T(8,128)(2,1)} broadcast(%param_2.4302), dimensions={2}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/...k,k->...k/dot_general" stack_frame_id=0} + %convert.1032 = f32[1,128,512]{1,2,0:T(8,128)} convert(%dot_general.916) + %param_1.5454 = bf16[1,128,576]{1,2,0:T(8,128)(2,1)S(1)} parameter(1) + %slice.1437 = bf16[1,128,512]{1,2,0:T(8,128)(2,1)} slice(%param_1.5454), slice={[0:1], [0:128], [0:512]}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/split" stack_frame_id=0} + %convert_element_type.3014 = f32[1,128,512]{1,2,0:T(8,128)} convert(%slice.1437), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %param_0.4851 = f32[128]{0:T(128)S(1)} parameter(0) + %mul.5142 = f32[1,128,512]{1,2,0:T(8,128)} broadcast(%param_0.4851), dimensions={1}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %mul.5141 = f32[1,128,512]{1,2,0:T(8,128)} multiply(%convert_element_type.3014, %mul.5142), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %convert_element_type.3013 = bf16[1,128,512]{1,2,0:T(8,128)(2,1)} convert(%mul.5141), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %convert.1033 = f32[1,128,512]{1,2,0:T(8,128)} convert(%convert_element_type.3013) + %dot_general.915 = f32[1,128,512]{1,2,0:T(8,128)} multiply(%convert.1032, %convert.1033), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/...k,k->...k/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","512"],"tuple_shapes":[],"layout":{"minor_to_major":["1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %convert.1034 = bf16[1,128,512]{1,2,0:T(8,128)(2,1)} convert(%dot_general.915) + ROOT %bitcast.2539 = bf16[128,512,1]{0,1,2:T(8,128)(2,1)} bitcast(%convert.1034) +} + +%fused_computation.101.clone.clone.clone.1 (param_0.4850: bf16[1,512,128,256]) -> bf16[512,128,256] { + %param_0.4850 = bf16[1,512,128,256]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.2538 = bf16[512,128,256]{2,0,1:T(8,128)(2,1)} bitcast(%param_0.4850), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} +} + +%fused_computation.130.clone.1 (param_0.4852: bf16[1,512,128,256], param_1.5455: f32[128], param_2.4303: bf16[1,128,576], param_3.2898: bf16[512]) -> bf16[1,128,128,256] { + %param_1.5455 = f32[128]{0:T(128)S(1)} parameter(1) + %param_2.4303 = bf16[1,128,576]{1,2,0:T(8,128)(2,1)S(1)} parameter(2) + %param_3.2898 = bf16[512]{0:T(512)(128)(2,1)S(1)} parameter(3) + %fusion.1119 = bf16[128,512,1]{0,1,2:T(8,128)(2,1)} fusion(%param_1.5455, %param_2.4303, %param_3.2898), kind=kLoop, calls=%fused_computation.168.clone.clone.1 + %param_0.4852 = bf16[1,512,128,256]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %fusion.1118 = bf16[512,128,256]{2,0,1:T(8,128)(2,1)} fusion(%param_0.4852), kind=kLoop, calls=%fused_computation.101.clone.clone.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} + %convolution.333 = bf16[128,128,256]{2,0,1:T(8,128)(2,1)} convolution(%fusion.1119, %fusion.1118), window={size=128 pad=127_127 rhs_reversal=1}, dim_labels=bf0_i0o->b0f, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} + ROOT %bitcast.2540 = bf16[1,128,128,256]{3,1,2,0:T(8,128)(2,1)S(1)} bitcast(%convolution.333), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} +} + +%fused_computation.132.clone.1 (param_0.4853: bf16[1,128,128,256]) -> (bf16[128,128,128], bf16[1,128,128,128]) { + %param_0.4853 = bf16[1,128,128,256]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %slice.1438 = bf16[1,128,128,128]{3,1,2,0:T(8,128)(2,1)} slice(%param_0.4853), slice={[0:1], [0:128], [0:128], [128:256]}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/split" stack_frame_id=0} + %bitcast.2541 = bf16[128,128,128]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%slice.1438), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/vmap(static_sparsity_splash)/jit(_splash_attention)/splash_mha_fwd_segmented_residuals/splash_mha_fwd_segmented_residuals/squeeze" stack_frame_id=0} + %slice.1439 = bf16[1,128,128,128]{3,1,2,0:T(8,128)(2,1)S(1)} slice(%param_0.4853), slice={[0:1], [0:128], [0:128], [0:128]}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/split" stack_frame_id=0} + ROOT %tuple.863 = (bf16[128,128,128]{2,1,0:T(8,128)(2,1)S(1)}, bf16[1,128,128,128]{3,1,2,0:T(8,128)(2,1)S(1)}) tuple(%bitcast.2541, %slice.1439) +} + +%fused_computation.120.clone.1 (param_0.4857: bf16[1,128,128,64], param_1.5458: bf16[1,128,128,128]) -> bf16[128,128,192] { + %param_1.5458 = bf16[1,128,128,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) + %constant.4258 = bf16[]{:T(256)} constant(-inf) + %pad.460 = bf16[1,128,128,192]{3,1,2,0:T(8,128)(2,1)} pad(%param_1.5458, %constant.4258), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/concatenate" stack_frame_id=0} + %convert.1059 = f32[1,128,128,192]{3,1,2,0:T(8,128)} convert(%pad.460) + %param_0.4857 = bf16[1,128,128,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %pad.459 = bf16[1,128,128,192]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.4857, %constant.4258), padding=0_0x0_0x0_0x128_0, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/concatenate" stack_frame_id=0} + %convert.1060 = f32[1,128,128,192]{3,1,2,0:T(8,128)} convert(%pad.459) + %maximum.98 = f32[1,128,128,192]{3,1,2,0:T(8,128)} maximum(%convert.1059, %convert.1060), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/concatenate" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","128","192"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %convert.1061 = bf16[1,128,128,192]{3,1,2,0:T(8,128)(2,1)} convert(%maximum.98) + ROOT %bitcast.2547 = bf16[128,128,192]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%convert.1061), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/vmap(static_sparsity_splash)/jit(_splash_attention)/splash_mha_fwd_segmented_residuals/splash_mha_fwd_segmented_residuals/squeeze" stack_frame_id=0} +} + +%fused_computation.133.clone.clone.clone.clone.clone.clone.clone.1 (param_0.4872: bf16[128,128,128]) -> bf16[128,128,128] { + %param_0.4872 = bf16[128,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.2557 = bf16[128,128,128]{2,0,1:T(8,128)(2,1)} bitcast(%param_0.4872) +} + +%fused_computation.109.clone.clone.clone.clone.clone.clone.clone.1 (param_0.4871: bf16[1,128,128,512]) -> bf16[128,128,512] { + %param_0.4871 = bf16[1,128,128,512]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.2556 = bf16[128,128,512]{2,1,0:T(8,128)(2,1)} bitcast(%param_0.4871), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} +} + +%region_5.7 (reduce_sum.31: f32[], reduce_sum.32: f32[]) -> f32[] { + %reduce_sum.31 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/reduce_sum"} + %reduce_sum.32 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/reduce_sum"} + ROOT %reduce_sum.33 = f32[]{:T(128)} add(%reduce_sum.31, %reduce_sum.32), metadata={op_name="reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%fused_computation.162.clone.1 (param_0.4873: bf16[1,128,512], param_1.5469: bf16[1,128,128,512], param_2.4311: bf16[128,128,128]) -> (f32[128], bf16[1,128,512]) { + %param_0.4873 = bf16[1,128,512]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert.1062 = f32[1,128,512]{2,1,0:T(8,128)} convert(%param_0.4873) + %param_2.4311 = bf16[128,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %fusion.112.clone.3 = bf16[128,128,128]{2,0,1:T(8,128)(2,1)} fusion(%param_2.4311), kind=kLoop, calls=%fused_computation.133.clone.clone.clone.clone.clone.clone.clone.1 + %param_1.5469 = bf16[1,128,128,512]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %fusion.111.clone.3 = bf16[128,128,512]{2,1,0:T(8,128)(2,1)} fusion(%param_1.5469), kind=kLoop, calls=%fused_computation.109.clone.clone.clone.clone.clone.clone.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} + %convolution.108.clone.3 = bf16[128,512,1]{1,0,2:T(8,128)(2,1)} convolution(%fusion.112.clone.3, %fusion.111.clone.3), window={size=128}, dim_labels=b0f_0io->bf0, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} + %bitcast.736.clone.3 = bf16[1,128,512]{2,1,0:T(8,128)(2,1)} bitcast(%convolution.108.clone.3), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} + %convert.1063 = f32[1,128,512]{2,1,0:T(8,128)} convert(%bitcast.736.clone.3) + %add.3136.clone.3 = f32[1,128,512]{2,1,0:T(8,128)} add(%convert.1062, %convert.1063), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/add" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","512"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %convert_element_type.3024 = f32[1,128,512]{2,1,0:T(8,128)} convert(%add.3136.clone.3), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %square.702 = f32[1,128,512]{2,1,0:T(8,128)} multiply(%convert_element_type.3024, %convert_element_type.3024), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/square" stack_frame_id=0} + %constant.1493.clone.24 = f32[]{:T(128)} constant(0) + %reduce.883 = f32[128]{0:T(128)S(1)} reduce(%square.702, %constant.1493.clone.24), dimensions={0,2}, to_apply=%region_5.7, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/reduce_sum" stack_frame_id=0} + %convert.1064 = bf16[1,128,512]{2,1,0:T(8,128)(2,1)S(1)} convert(%add.3136.clone.3) + ROOT %tuple.869 = (f32[128]{0:T(128)S(1)}, bf16[1,128,512]{2,1,0:T(8,128)(2,1)S(1)}) tuple(%reduce.883, %convert.1064) +} + +%fused_computation.115.clone.1 (param_0.4880: bf16[3,18432,128], param_1.5474: s32[]) -> bf16[1,18432,128] { + %param_0.4880 = bf16[3,18432,128]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.5474 = s32[]{:T(128)S(6)} parameter(1) + %constant.1492.clone.53 = s32[]{:T(128)} constant(0) + ROOT %dynamic-slice.535 = bf16[1,18432,128]{2,1,0:T(8,128)(2,1)S(1)} dynamic-slice(%param_0.4880, %param_1.5474, %constant.1492.clone.53, %constant.1492.clone.53), dynamic_slice_sizes={1,18432,128}, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_index_in_dim)/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} +} + +%fused_computation.116.clone.1 (param_0.4879: bf16[3,128,18432], param_1.5473: s32[]) -> bf16[1,128,18432] { + %param_0.4879 = bf16[3,128,18432]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.5473 = s32[]{:T(128)S(6)} parameter(1) + %constant.1492.clone.52 = s32[]{:T(128)} constant(0) + ROOT %dynamic-slice.534 = bf16[1,128,18432]{2,1,0:T(8,128)(2,1)S(1)} dynamic-slice(%param_0.4879, %param_1.5473, %constant.1492.clone.52, %constant.1492.clone.52), dynamic_slice_sizes={1,128,18432}, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_index_in_dim)/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} +} + +%fused_computation.117.clone.1 (param_0.4875: bf16[3,128,18432], param_1.5470: s32[]) -> bf16[1,128,18432] { + %param_0.4875 = bf16[3,128,18432]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.5470 = s32[]{:T(128)S(6)} parameter(1) + %constant.1492.clone.51 = s32[]{:T(128)} constant(0) + ROOT %dynamic-slice.533 = bf16[1,128,18432]{2,1,0:T(8,128)(2,1)S(1)} dynamic-slice(%param_0.4875, %param_1.5470, %constant.1492.clone.51, %constant.1492.clone.51), dynamic_slice_sizes={1,128,18432}, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_index_in_dim)/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} +} + +%fused_computation.202.clone.1 (param_0.4874: f32[128]) -> f32[128] { + %param_0.4874 = f32[128]{0:T(128)S(1)} parameter(0) + %constant.1494.clone.13 = f32[]{:T(128)} constant(0.001953125) + %broadcast.4713 = f32[128]{0:T(128)} broadcast(%constant.1494.clone.13), dimensions={}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call" stack_frame_id=0} + %div.2873 = f32[128]{0:T(128)} multiply(%param_0.4874, %broadcast.4713), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/div" stack_frame_id=0} + %constant.1495.clone.16 = f32[]{:T(128)} constant(1e-06) + %broadcast.4712 = f32[128]{0:T(128)} broadcast(%constant.1495.clone.16), dimensions={}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call" stack_frame_id=0} + %add.3972 = f32[128]{0:T(128)} add(%div.2873, %broadcast.4712), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/add" stack_frame_id=0} + ROOT %rsqrt.369 = f32[128]{0:T(128)S(1)} rsqrt(%add.3972), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/rsqrt" stack_frame_id=0} +} + +%fused_computation.166.clone.1.clone.1 (param_0.4877: bf16[1,128,512], param_1.5471: f32[128], param_2.4312: bf16[512]) -> bf16[128,512] { + %param_2.4312 = bf16[512]{0:T(512)(128)(2,1)S(1)} parameter(2) + %dot_general.922 = bf16[1,128,512]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.4312), dimensions={2}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/...k,k->...k/dot_general" stack_frame_id=0} + %convert.1041 = f32[1,128,512]{2,1,0:T(8,128)} convert(%dot_general.922) + %param_0.4877 = bf16[1,128,512]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.3026 = f32[1,128,512]{2,1,0:T(8,128)} convert(%param_0.4877), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %param_1.5471 = f32[128]{0:T(128)S(1)} parameter(1) + %mul.5158 = f32[1,128,512]{2,1,0:T(8,128)} broadcast(%param_1.5471), dimensions={1}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %mul.5157 = f32[1,128,512]{2,1,0:T(8,128)} multiply(%convert_element_type.3026, %mul.5158), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %convert_element_type.3025 = bf16[1,128,512]{2,1,0:T(8,128)(2,1)} convert(%mul.5157), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %convert.1042 = f32[1,128,512]{2,1,0:T(8,128)} convert(%convert_element_type.3025) + %dot_general.921 = f32[1,128,512]{2,1,0:T(8,128)} multiply(%convert.1041, %convert.1042), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/...k,k->...k/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","512"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %convert.1043 = bf16[1,128,512]{2,1,0:T(8,128)(2,1)} convert(%dot_general.921) + ROOT %bitcast.2559 = bf16[128,512]{1,0:T(8,128)(2,1)} bitcast(%convert.1043) +} + +%fused_computation.107.clone.clone.clone.1 (param_0.4876: bf16[1,512,18432]) -> bf16[512,18432] { + %param_0.4876 = bf16[1,512,18432]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.2558 = bf16[512,18432]{1,0:T(8,128)(2,1)} bitcast(%param_0.4876), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} +} + +%fused_computation.111.clone.1 (param_0.4878: bf16[1,512,18432], param_1.5472: bf16[1,128,512], param_2.4313: f32[128], param_3.2903: bf16[512]) -> bf16[1,128,18432] { + %param_1.5472 = bf16[1,128,512]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %param_2.4313 = f32[128]{0:T(128)S(1)} parameter(2) + %param_3.2903 = bf16[512]{0:T(512)(128)(2,1)S(1)} parameter(3) + %fusion.1123 = bf16[128,512]{1,0:T(8,128)(2,1)} fusion(%param_1.5472, %param_2.4313, %param_3.2903), kind=kLoop, calls=%fused_computation.166.clone.1.clone.1 + %param_0.4878 = bf16[1,512,18432]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %fusion.1122 = bf16[512,18432]{1,0:T(8,128)(2,1)} fusion(%param_0.4878), kind=kLoop, calls=%fused_computation.107.clone.clone.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} + %convolution.335 = bf16[128,18432]{1,0:T(8,128)(2,1)} convolution(%fusion.1123, %fusion.1122), dim_labels=bf_io->bf, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} + ROOT %bitcast.2560 = bf16[1,128,18432]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%convolution.335), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} +} + +%fused_computation.166.clone.clone.clone.1 (param_0.4883: bf16[1,128,512], param_1.5475: f32[128], param_2.4314: bf16[512]) -> bf16[128,512] { + %param_2.4314 = bf16[512]{0:T(512)(128)(2,1)S(1)} parameter(2) + %dot_general.924 = bf16[1,128,512]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.4314), dimensions={2}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/...k,k->...k/dot_general" stack_frame_id=0} + %convert.1044 = f32[1,128,512]{2,1,0:T(8,128)} convert(%dot_general.924) + %param_0.4883 = bf16[1,128,512]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.3028 = f32[1,128,512]{2,1,0:T(8,128)} convert(%param_0.4883), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %param_1.5475 = f32[128]{0:T(128)S(1)} parameter(1) + %mul.5160 = f32[1,128,512]{2,1,0:T(8,128)} broadcast(%param_1.5475), dimensions={1}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %mul.5159 = f32[1,128,512]{2,1,0:T(8,128)} multiply(%convert_element_type.3028, %mul.5160), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %convert_element_type.3027 = bf16[1,128,512]{2,1,0:T(8,128)(2,1)} convert(%mul.5159), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %convert.1045 = f32[1,128,512]{2,1,0:T(8,128)} convert(%convert_element_type.3027) + %dot_general.923 = f32[1,128,512]{2,1,0:T(8,128)} multiply(%convert.1044, %convert.1045), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/...k,k->...k/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","512"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %convert.1046 = bf16[1,128,512]{2,1,0:T(8,128)(2,1)} convert(%dot_general.923) + ROOT %bitcast.2563 = bf16[128,512]{1,0:T(8,128)(2,1)} bitcast(%convert.1046) +} + +%fused_computation.105.clone.clone.clone.clone.clone.clone.clone.clone.clone.clone.clone.clone.clone.1 (param_0.4882: bf16[1,512,18432]) -> bf16[512,18432] { + %param_0.4882 = bf16[1,512,18432]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.2562 = bf16[512,18432]{1,0:T(8,128)(2,1)} bitcast(%param_0.4882), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} +} + +%fused_computation.114.clone.clone.clone.clone.clone.clone.clone.1 (param_0.4884: bf16[1,512,18432], param_1.5476: bf16[1,128,18432], param_2.4315: bf16[1,128,512], param_3.2904: f32[128], param_4.2205: bf16[512]) -> bf16[128,18432] { + %param_1.5476 = bf16[1,128,18432]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %convert.1047 = f32[1,128,18432]{2,1,0:T(8,128)} convert(%param_1.5476) + %constant.1498.clone.10 = bf16[]{:T(256)} constant(1) + %jit_silu_.98 = bf16[1,128,18432]{2,1,0:T(8,128)(2,1)} broadcast(%constant.1498.clone.10), dimensions={}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/jit(silu)" stack_frame_id=0} + %convert.1048 = f32[1,128,18432]{2,1,0:T(8,128)} convert(%jit_silu_.98) + %neg.274 = f32[1,128,18432]{2,1,0:T(8,128)} negate(%convert.1047), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/jit(silu)/neg" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","18432"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %exp.567 = f32[1,128,18432]{2,1,0:T(8,128)} exponential(%neg.274), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/jit(silu)/exp" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","18432"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %add.3973 = f32[1,128,18432]{2,1,0:T(8,128)} add(%exp.567, %convert.1048), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/jit(silu)/add" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","18432"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %div.2874 = f32[1,128,18432]{2,1,0:T(8,128)} divide(%convert.1048, %add.3973), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/jit(silu)/div" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","18432"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %mul.5162 = f32[1,128,18432]{2,1,0:T(8,128)} multiply(%convert.1047, %div.2874), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/jit(silu)/mul" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","18432"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %param_2.4315 = bf16[1,128,512]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %param_3.2904 = f32[128]{0:T(128)S(1)} parameter(3) + %param_4.2205 = bf16[512]{0:T(512)(128)(2,1)S(1)} parameter(4) + %fusion.1125 = bf16[128,512]{1,0:T(8,128)(2,1)} fusion(%param_2.4315, %param_3.2904, %param_4.2205), kind=kLoop, calls=%fused_computation.166.clone.clone.clone.1 + %param_0.4884 = bf16[1,512,18432]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %fusion.1124 = bf16[512,18432]{1,0:T(8,128)(2,1)} fusion(%param_0.4884), kind=kLoop, calls=%fused_computation.105.clone.clone.clone.clone.clone.clone.clone.clone.clone.clone.clone.clone.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} + %convolution.336 = bf16[128,18432]{1,0:T(8,128)(2,1)} convolution(%fusion.1125, %fusion.1124), dim_labels=bf_io->bf, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} + %bitcast.2565 = bf16[1,128,18432]{2,1,0:T(8,128)(2,1)} bitcast(%convolution.336), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} + %convert.1049 = f32[1,128,18432]{2,1,0:T(8,128)} convert(%bitcast.2565) + %mul.5161 = f32[1,128,18432]{2,1,0:T(8,128)} multiply(%mul.5162, %convert.1049), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","18432"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %convert.1050 = bf16[1,128,18432]{2,1,0:T(8,128)(2,1)} convert(%mul.5161) + ROOT %bitcast.2564 = bf16[128,18432]{1,0:T(8,128)(2,1)} bitcast(%convert.1050) +} + +%fused_computation.103.clone.clone.clone.clone.clone.1 (param_0.4881: bf16[1,18432,512]) -> bf16[18432,512] { + %param_0.4881 = bf16[1,18432,512]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.2561 = bf16[18432,512]{1,0:T(8,128)(2,1)} bitcast(%param_0.4881), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} +} + +%fused_computation.146.clone.1 (param_0.4885: bf16[1,128,512], param_1.5477: bf16[1,18432,512], param_2.4316: bf16[1,512,18432], param_3.2905: bf16[1,128,18432], param_4.2206: f32[128], param_5.1916: bf16[512]) -> bf16[1,128,512] { + %param_2.4316 = bf16[1,512,18432]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %param_3.2905 = bf16[1,128,18432]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) + %param_0.4885 = bf16[1,128,512]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %param_4.2206 = f32[128]{0:T(128)S(1)} parameter(4) + %param_5.1916 = bf16[512]{0:T(512)(128)(2,1)S(1)} parameter(5) + %fusion.65.clone.7 = bf16[128,18432]{1,0:T(8,128)(2,1)} fusion(%param_2.4316, %param_3.2905, %param_0.4885, %param_4.2206, %param_5.1916), kind=kOutput, calls=%fused_computation.114.clone.clone.clone.clone.clone.clone.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} + %param_1.5477 = bf16[1,18432,512]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %fusion.1126 = bf16[18432,512]{1,0:T(8,128)(2,1)} fusion(%param_1.5477), kind=kLoop, calls=%fused_computation.103.clone.clone.clone.clone.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} + %convolution.337 = bf16[128,512]{1,0:T(8,128)(2,1)} convolution(%fusion.65.clone.7, %fusion.1126), dim_labels=bf_io->bf, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} + %bitcast.2566 = bf16[1,128,512]{2,1,0:T(8,128)(2,1)} bitcast(%convolution.337), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} + %convert.1065 = f32[1,128,512]{2,1,0:T(8,128)} convert(%bitcast.2566) + %convert.1066 = f32[1,128,512]{2,1,0:T(8,128)} convert(%param_0.4885) + %add.3974 = f32[1,128,512]{2,1,0:T(8,128)} add(%convert.1065, %convert.1066), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/add" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","512"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + ROOT %convert.1067 = bf16[1,128,512]{2,1,0:T(8,128)(2,1)S(1)} convert(%add.3974) +} + +%async_computation.48 (param_0.4952: s8[16777216]) -> s8[4194304] { + %param_0.4952 = s8[16777216]{0:T(1024)(128)(4,1)} parameter(0) + ROOT %slice.1490 = s8[4194304]{0:T(1024)(128)(4,1)S(1)} slice(%param_0.4952), slice={[0:4194304]}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%async_computation.49 (param_0.4953: s8[16777216]) -> s8[4194304] { + %param_0.4953 = s8[16777216]{0:T(1024)(128)(4,1)} parameter(0) + ROOT %slice.1492 = s8[4194304]{0:T(1024)(128)(4,1)S(1)} slice(%param_0.4953), slice={[4194304:8388608]}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%async_computation.50 (param_0.4954: s8[16777216]) -> s8[4194304] { + %param_0.4954 = s8[16777216]{0:T(1024)(128)(4,1)} parameter(0) + ROOT %slice.1494 = s8[4194304]{0:T(1024)(128)(4,1)S(1)} slice(%param_0.4954), slice={[8388608:12582912]}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%async_computation.51 (param_0.4955: s8[16777216]) -> s8[4194304] { + %param_0.4955 = s8[16777216]{0:T(1024)(128)(4,1)} parameter(0) + ROOT %slice.1496 = s8[4194304]{0:T(1024)(128)(4,1)S(1)} slice(%param_0.4955), slice={[12582912:16777216]}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%wide.region_1.8_spmd.sunk.clone.clone.clone.sunk (wide.param.5: (s32[], bf16[1,128,512], u32[3], u32[3,4], u32[3], /*index=5*/u32[3,4], u32[3], u32[3,4], u32[3], u32[3,4], /*index=10*/u32[3], u32[3,4], u32[3], u32[3,4], u32[3], /*index=15*/u32[3,4], u32[3], u32[3,4], u32[3], u32[3,4], /*index=20*/bf16[3,1,128,512], bf16[3,512], bf16[3,384], bf16[3,512], bf16[3,128,1536], /*index=25*/bf16[3,384,128,192], bf16[3,128], bf16[3,128,576], bf16[3,128,128,256], f32[1,128,1,32], /*index=30*/f32[1,128,1,32], bf16[3,128,128,128], bf16[3,128,18432], bf16[3,128,18432], bf16[3,18432,128], /*index=35*/u32[3], u32[3,4], u32[3], u32[3,4], u32[3], /*index=40*/u32[3,4], u32[3], u32[3,4], u32[3], u32[3,4], /*index=45*/u32[3], u32[3,4], u32[3], u32[3,4], u32[3], /*index=50*/u32[3,4], u32[3], u32[3,4], f32[128,32], f32[128,32], /*index=55*/s32[128], s32[], s8[1,1,1], s8[1,1,1], bf16[])) -> (s32[], bf16[1,128,512], u32[3], u32[3,4], u32[3], /*index=5*/u32[3,4], u32[3], u32[3,4], u32[3], u32[3,4], /*index=10*/u32[3], u32[3,4], u32[3], u32[3,4], u32[3], /*index=15*/u32[3,4], u32[3], u32[3,4], u32[3], u32[3,4], /*index=20*/bf16[3,1,128,512], bf16[3,512], bf16[3,384], bf16[3,512], bf16[3,128,1536], /*index=25*/bf16[3,384,128,192], bf16[3,128], bf16[3,128,576], bf16[3,128,128,256], f32[1,128,1,32], /*index=30*/f32[1,128,1,32], bf16[3,128,128,128], bf16[3,128,18432], bf16[3,128,18432], bf16[3,18432,128], /*index=35*/u32[3], u32[3,4], u32[3], u32[3,4], u32[3], /*index=40*/u32[3,4], u32[3], u32[3,4], u32[3], u32[3,4], /*index=45*/u32[3], u32[3,4], u32[3], u32[3,4], u32[3], /*index=50*/u32[3,4], u32[3], u32[3,4], f32[128,32], f32[128,32], /*index=55*/s32[128], s32[], s8[1,1,1], s8[1,1,1], bf16[]) { + %constant.1752.clone..sunk.7 = s32[]{:T(128)} constant(1) + %constant.4192..sunk.1 = bf16[]{:T(256)} constant(-0), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %wide.param.5 = (s32[]{:T(128)}, bf16[1,128,512]{2,1,0:T(8,128)(2,1)S(1)}, u32[3]{0:T(128)}, u32[3,4]{1,0:T(4,128)}, u32[3]{0:T(128)}, /*index=5*/u32[3,4]{1,0:T(4,128)}, u32[3]{0:T(128)}, u32[3,4]{1,0:T(4,128)}, u32[3]{0:T(128)}, u32[3,4]{1,0:T(4,128)}, /*index=10*/u32[3]{0:T(128)}, u32[3,4]{1,0:T(4,128)}, u32[3]{0:T(128)}, u32[3,4]{1,0:T(4,128)}, u32[3]{0:T(128)}, /*index=15*/u32[3,4]{1,0:T(4,128)}, u32[3]{0:T(128)}, u32[3,4]{1,0:T(4,128)}, u32[3]{0:T(128)}, u32[3,4]{1,0:T(4,128)}, /*index=20*/bf16[3,1,128,512]{3,2,1,0:T(8,128)(2,1)}, bf16[3,512]{1,0:T(4,128)(2,1)}, bf16[3,384]{1,0:T(4,128)(2,1)}, bf16[3,512]{1,0:T(4,128)(2,1)}, bf16[3,128,1536]{2,1,0:T(8,128)(2,1)}, /*index=25*/bf16[3,384,128,192]{2,1,3,0:T(8,128)(2,1)}, bf16[3,128]{1,0:T(4,128)(2,1)}, bf16[3,128,576]{2,1,0:T(8,128)(2,1)S(1)}, bf16[3,128,128,256]{3,1,2,0:T(8,128)(2,1)}, f32[1,128,1,32]{3,1,2,0:T(8,128)S(1)}, /*index=30*/f32[1,128,1,32]{3,1,2,0:T(8,128)S(1)}, bf16[3,128,128,128]{3,2,1,0:T(8,128)(2,1)}, bf16[3,128,18432]{2,1,0:T(8,128)(2,1)}, bf16[3,128,18432]{2,1,0:T(8,128)(2,1)}, bf16[3,18432,128]{2,1,0:T(8,128)(2,1)}, /*index=35*/u32[3]{0:T(128)}, u32[3,4]{1,0:T(4,128)}, u32[3]{0:T(128)}, u32[3,4]{1,0:T(4,128)}, u32[3]{0:T(128)}, /*index=40*/u32[3,4]{1,0:T(4,128)}, u32[3]{0:T(128)}, u32[3,4]{1,0:T(4,128)}, u32[3]{0:T(128)}, u32[3,4]{1,0:T(4,128)}, /*index=45*/u32[3]{0:T(128)}, u32[3,4]{1,0:T(4,128)}, u32[3]{0:T(128)}, u32[3,4]{1,0:T(4,128)}, u32[3]{0:T(128)}, /*index=50*/u32[3,4]{1,0:T(4,128)}, u32[3]{0:T(128)}, u32[3,4]{1,0:T(4,128)}, f32[128,32]{1,0:T(8,128)S(1)}, f32[128,32]{1,0:T(8,128)S(1)}, /*index=55*/s32[128]{0:T(128)}, s32[]{:T(128)}, s8[1,1,1]{2,1,0:T(4,128)(4,1)}, s8[1,1,1]{2,1,0:T(4,128)(4,1)}, bf16[]{:T(256)}) parameter(0) + %get-tuple-element.8821 = s32[]{:T(128)} get-tuple-element(%wide.param.5), index=0 + %copy.1644 = s32[]{:T(128)S(6)} copy(%get-tuple-element.8821), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %add.3975 = s32[]{:T(128)} add(%copy.1644, %constant.1752.clone..sunk.7), metadata={op_name="jit(train_step)/jvp()/while/body/add" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %get-tuple-element.8841 = bf16[3,1,128,512]{3,2,1,0:T(8,128)(2,1)} get-tuple-element(%wide.param.5), index=20 + %get-tuple-element.8822 = bf16[1,128,512]{2,1,0:T(8,128)(2,1)S(1)} get-tuple-element(%wide.param.5), index=1 + %bitcast_dynamic-update-slice_fusion.17 = bf16[3,1,128,512]{3,2,1,0:T(8,128)(2,1)} fusion(%get-tuple-element.8841, %copy.1644, %get-tuple-element.8822), kind=kLoop, calls=%fused_computation.171.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_update_index_in_dim)/dynamic_update_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[{"indices":["0","3"]}]}} + %get-tuple-element.8912 = bf16[3,128,128,128]{3,2,1,0:T(8,128)(2,1)} get-tuple-element(%wide.param.5), index=31 + %dynamic-slice_convert_fusion.40 = bf16[1,128,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} fusion(%get-tuple-element.8912, %copy.1644), kind=kLoop, calls=%fused_computation.129.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %all-gather.340 = bf16[1,128,128,512]{3,2,1,0:T(8,128)(2,1)} all-gather(%dynamic-slice_convert_fusion.40), channel_id=288, replica_groups=mesh['axis_0'=1,'axis_1'=1,'axis_2'=4] {'axis_2'}, dimensions={3}, use_global_device_ids=true, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"barrier_config":{"barrier_type":"CUSTOM","id":"0"},"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %get-tuple-element.8938 = s8[1,1,1]{2,1,0:T(4,128)(4,1)} get-tuple-element(%wide.param.5), index=57, metadata={op_name="jit(train_step)/sharding_constraint" stack_frame_id=0} + %get-tuple-element.8939 = s8[1,1,1]{2,1,0:T(4,128)(4,1)} get-tuple-element(%wide.param.5), index=58, metadata={op_name="jit(train_step)/sharding_constraint" stack_frame_id=0} + %get-tuple-element.8906 = bf16[3,384,128,192]{2,1,3,0:T(8,128)(2,1)} get-tuple-element(%wide.param.5), index=25 + %dynamic-slice_convert_fusion.41 = bf16[1,384,128,192]{2,1,3,0:T(8,128)(2,1)S(1)} fusion(%get-tuple-element.8906, %copy.1644), kind=kLoop, calls=%fused_computation.99.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %all-gather.341 = bf16[1,1536,128,192]{2,1,3,0:T(8,128)(2,1)} all-gather(%dynamic-slice_convert_fusion.41), channel_id=284, replica_groups=mesh['axis_0'=4,'axis_1'=1,'axis_2'=1] {'axis_0'}, dimensions={1}, use_global_device_ids=true, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"barrier_config":{"barrier_type":"CUSTOM","id":"0"},"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %bitcast.2567 = bf16[1536,128,192]{1,0,2:T(8,128)(2,1)} bitcast(%all-gather.341), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} + %copy.1595 = bf16[1536,128,192]{0,2,1:T(8,128)(2,1)} copy(%bitcast.2567), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %get-tuple-element.8905 = bf16[3,128,1536]{2,1,0:T(8,128)(2,1)} get-tuple-element(%wide.param.5), index=24 + %dynamic-slice_convert_fusion.42 = bf16[1,128,1536]{2,1,0:T(8,128)(2,1)S(1)} fusion(%get-tuple-element.8905, %copy.1644), kind=kLoop, calls=%fused_computation.140.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %all-gather.342 = bf16[1,512,1536]{2,1,0:T(8,128)(2,1)S(1)} all-gather(%dynamic-slice_convert_fusion.42), channel_id=283, replica_groups=mesh['axis_0'=4,'axis_1'=1] {'axis_0'}, dimensions={1}, use_global_device_ids=true, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"barrier_config":{"barrier_type":"CUSTOM","id":"0"},"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %fusion.1127 = f32[128]{0:T(128)S(1)} fusion(%get-tuple-element.8822), kind=kLoop, calls=%fused_computation.158.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %add_rsqrt_fusion.11 = f32[128]{0:T(128)S(1)} fusion(%fusion.1127), kind=kLoop, calls=%fused_computation.205.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/rsqrt" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %get-tuple-element.8902 = bf16[3,512]{1,0:T(4,128)(2,1)} get-tuple-element(%wide.param.5), index=21 + %get-tuple-element.8904 = bf16[3,512]{1,0:T(4,128)(2,1)} get-tuple-element(%wide.param.5), index=23 + %fusion.1128 = (bf16[512]{0:T(512)(128)(2,1)S(1)}, bf16[512]{0:T(512)(128)(2,1)S(1)}) fusion(%get-tuple-element.8902, %copy.1644, %get-tuple-element.8904), kind=kLoop, calls=%fused_computation.197.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"fusion_formation_source":{"fusion_creating_compiler_pass":"TpuAdvancedMOF","fusion_id":"42"},"aliasing_operands":{"lists":[]}} + %get-tuple-element.8744 = bf16[512]{0:T(512)(128)(2,1)S(1)} get-tuple-element(%fusion.1128), index=1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %multiply_reduce_fusion.96 = (f32[128]{0:T(128)S(1)}, f32[1,128,1536]{2,1,0:T(8,128)S(1)}) fusion(%all-gather.342, %get-tuple-element.8822, %add_rsqrt_fusion.11, %get-tuple-element.8744), kind=kOutput, calls=%fused_computation.139.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %get-tuple-element.8745 = f32[1,128,1536]{2,1,0:T(8,128)S(1)} get-tuple-element(%multiply_reduce_fusion.96), index=1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/reduce_sum" stack_frame_id=0} + %get-tuple-element.8746 = f32[128]{0:T(128)S(1)} get-tuple-element(%multiply_reduce_fusion.96), index=0, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/reduce_sum" stack_frame_id=0} + %add_rsqrt_fusion.12 = f32[128]{0:T(128)S(1)} fusion(%get-tuple-element.8746), kind=kLoop, calls=%fused_computation.204.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/rsqrt" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %get-tuple-element.8903 = bf16[3,384]{1,0:T(4,128)(2,1)} get-tuple-element(%wide.param.5), index=22 + %fusion.1129 = bf16[1,384]{1,0:T(2,128)(2,1)S(1)} fusion(%get-tuple-element.8903, %copy.1644), kind=kLoop, calls=%fused_computation.200.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %all-gather.343 = bf16[1,1536]{1,0:T(2,128)(2,1)S(1)} all-gather(%fusion.1129), channel_id=282, replica_groups=mesh['axis_0'=4] {'axis_0'}, dimensions={1}, use_global_device_ids=true, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/...k,k->...k/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"barrier_config":{"barrier_type":"CUSTOM","id":"0"},"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %reduce.884 = bf16[1536]{0:T(1024)(128)(2,1)S(1)} reduce(%all-gather.343, %constant.4192..sunk.1), dimensions={0}, to_apply=%all-gather.290.reduce_sub_computation, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/...k,k->...k/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %convolution_bitcast_fusion.10 = bf16[1,128,128,192]{3,1,2,0:T(8,128)(2,1)S(1)} fusion(%copy.1595, %get-tuple-element.8745, %add_rsqrt_fusion.12, %reduce.884), kind=kOutput, calls=%fused_computation.131.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %slice.1444 = bf16[1,128,128,64]{3,1,2,0:T(8,128)(2,1)S(1)} slice(%convolution_bitcast_fusion.10), slice={[0:1], [0:128], [0:128], [128:192]}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/split" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %copy.1596 = bf16[1,128,128,64]{2,1,3,0:T(8,128)(2,1)S(1)} copy(%slice.1444), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/split" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %bitcast.2611 = bf16[1,128,128,32,2]{2,1,4,3,0:T(8,128)(2,1)S(1)} bitcast(%copy.1596) + %fusion.1130 = (bf16[1,128,128,32,1]{2,1,4,3,0:T(8,128)(2,1)S(1)}, bf16[1,128,128,32,1]{2,1,4,3,0:T(8,128)(2,1)S(1)}) fusion(%bitcast.2611), kind=kLoop, calls=%fused_computation.1415.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %get-tuple-element.8747 = bf16[1,128,128,32,1]{2,1,4,3,0:T(8,128)(2,1)S(1)} get-tuple-element(%fusion.1130), index=1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/slice" stack_frame_id=0} + %copy.1597 = bf16[1,128,128,32,1]{3,1,2,4,0:T(8,128)(2,1)S(1)} copy(%get-tuple-element.8747), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %get-tuple-element.8748 = bf16[1,128,128,32,1]{2,1,4,3,0:T(8,128)(2,1)S(1)} get-tuple-element(%fusion.1130), index=0, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/slice" stack_frame_id=0} + %copy.1598 = bf16[1,128,128,32,1]{3,1,2,4,0:T(8,128)(2,1)S(1)} copy(%get-tuple-element.8748), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %get-tuple-element.8935 = f32[128,32]{1,0:T(8,128)S(1)} get-tuple-element(%wide.param.5), index=54 + %get-tuple-element.8934 = f32[128,32]{1,0:T(8,128)S(1)} get-tuple-element(%wide.param.5), index=53 + %fusion.1131 = (bf16[1,128,128,32]{3,1,2,0:T(8,128)(2,1)S(1)}, bf16[1,128,128,32]{3,1,2,0:T(8,128)(2,1)S(1)}) fusion(%copy.1597, %copy.1598, %get-tuple-element.8935, %get-tuple-element.8934), kind=kLoop, calls=%fused_computation.126.clone.1, metadata={op_name="convert.317"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %get-tuple-element.8749 = bf16[1,128,128,32]{3,1,2,0:T(8,128)(2,1)S(1)} get-tuple-element(%fusion.1131), index=1, metadata={op_name="convert.317"} + %get-tuple-element.8750 = bf16[1,128,128,32]{3,1,2,0:T(8,128)(2,1)S(1)} get-tuple-element(%fusion.1131), index=0, metadata={op_name="convert.317"} + %slice.1445 = bf16[1,128,128,128]{3,1,2,0:T(8,128)(2,1)S(1)} slice(%convolution_bitcast_fusion.10), slice={[0:1], [0:128], [0:128], [0:128]}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/split" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %fusion.1132 = bf16[128,128,192]{2,1,0:T(8,128)(2,1)S(1)} fusion(%get-tuple-element.8749, %get-tuple-element.8750, %slice.1445), kind=kLoop, calls=%fused_computation.119.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/vmap(static_sparsity_splash)/jit(_splash_attention)/splash_mha_fwd_segmented_residuals/splash_mha_fwd_segmented_residuals/squeeze" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %get-tuple-element.8908 = bf16[3,128,576]{2,1,0:T(8,128)(2,1)S(1)} get-tuple-element(%wide.param.5), index=27 + %dynamic-slice_convert_fusion.43 = bf16[1,128,576]{2,1,0:T(8,128)(2,1)S(1)} fusion(%get-tuple-element.8908, %copy.1644), kind=kLoop, calls=%fused_computation.155.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %all-gather.344 = bf16[1,512,576]{2,1,0:T(8,128)(2,1)S(1)} all-gather(%dynamic-slice_convert_fusion.43), channel_id=286, replica_groups=mesh['axis_0'=4,'axis_1'=1] {'axis_0'}, dimensions={1}, use_global_device_ids=true, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"barrier_config":{"barrier_type":"CUSTOM","id":"0"},"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %fusion.1133 = bf16[1,128,576]{1,2,0:T(8,128)(2,1)S(1)} fusion(%all-gather.344, %get-tuple-element.8822, %add_rsqrt_fusion.11, %get-tuple-element.8744), kind=kOutput, calls=%fused_computation.143.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %slice.1446 = bf16[1,128,64]{1,2,0:T(8,128)(2,1)S(1)} slice(%fusion.1133), slice={[0:1], [0:128], [512:576]}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/split" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %bitcast.2682 = bf16[1,128,1,32,2]{1,4,3,2,0:T(2,128)(2,1)S(1)} bitcast(%slice.1446) + %fusion.1134 = (bf16[1,128,1,32,1]{1,4,3,2,0:T(2,128)(2,1)S(1)}, bf16[1,128,1,32,1]{1,4,3,2,0:T(2,128)(2,1)S(1)}) fusion(%bitcast.2682), kind=kLoop, calls=%fused_computation.1416.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %get-tuple-element.8751 = bf16[1,128,1,32,1]{1,4,3,2,0:T(2,128)(2,1)S(1)} get-tuple-element(%fusion.1134), index=1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/slice" stack_frame_id=0} + %copy.1599 = bf16[1,128,1,32,1]{3,1,4,2,0:T(8,128)(2,1)S(1)} copy(%get-tuple-element.8751), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %get-tuple-element.8752 = bf16[1,128,1,32,1]{1,4,3,2,0:T(2,128)(2,1)S(1)} get-tuple-element(%fusion.1134), index=0, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/slice" stack_frame_id=0} + %copy.1600 = bf16[1,128,1,32,1]{3,1,4,2,0:T(8,128)(2,1)S(1)} copy(%get-tuple-element.8752), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %bitcast.2669 = f32[1,128,1,32]{3,1,2,0:T(8,128)S(1)} bitcast(%get-tuple-element.8935) + %bitcast.2666 = f32[1,128,1,32]{3,1,2,0:T(8,128)S(1)} bitcast(%get-tuple-element.8934) + %fusion.1135 = (bf16[1,128,32]{2,1,0:T(8,128)(2,1)S(1)}, bf16[1,128,32]{2,1,0:T(8,128)(2,1)S(1)}) fusion(%bitcast.2669, %bitcast.2666, %copy.1599, %copy.1600), kind=kLoop, calls=%fused_computation.179.clone.1, metadata={op_name="convert.315"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %get-tuple-element.8753 = bf16[1,128,32]{2,1,0:T(8,128)(2,1)S(1)} get-tuple-element(%fusion.1135), index=1, metadata={op_name="convert.315"} + %get-tuple-element.8754 = bf16[1,128,32]{2,1,0:T(8,128)(2,1)S(1)} get-tuple-element(%fusion.1135), index=0, metadata={op_name="convert.315"} + %maximum_bitcast_fusion.27 = bf16[128,64]{1,0:T(8,128)(2,1)S(1)} fusion(%get-tuple-element.8753, %get-tuple-element.8754), kind=kLoop, calls=%fused_computation.191.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/concatenate" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %broadcast_in_dim.1649 = bf16[1,128,128,64]{3,1,2,0:T(8,128)(2,1)S(1)} broadcast(%maximum_bitcast_fusion.27), dimensions={1,3}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/broadcast_in_dim" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"output_chunk_bound_config":{"output_chunk_bound":["8","128"]},"aliasing_operands":{"lists":[]}} + %get-tuple-element.8909 = bf16[3,128,128,256]{3,1,2,0:T(8,128)(2,1)} get-tuple-element(%wide.param.5), index=28 + %dynamic-slice_convert_fusion.44 = bf16[1,128,128,256]{3,1,2,0:T(8,128)(2,1)S(1)} fusion(%get-tuple-element.8909, %copy.1644), kind=kLoop, calls=%fused_computation.110.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %get-tuple-element.8907 = bf16[3,128]{1,0:T(4,128)(2,1)} get-tuple-element(%wide.param.5), index=26 + %copy-start.242 = (bf16[3,128]{1,0:T(4,128)(2,1)S(1)}, bf16[3,128]{1,0:T(4,128)(2,1)}, u32[]{:S(2)}) copy-start(%get-tuple-element.8907) + %all-gather.345 = bf16[1,512,128,256]{3,1,2,0:T(8,128)(2,1)S(1)} all-gather(%dynamic-slice_convert_fusion.44), channel_id=287, replica_groups=mesh['axis_0'=4,'axis_1'=1,'axis_2'=1] {'axis_0'}, dimensions={1}, use_global_device_ids=true, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"barrier_config":{"barrier_type":"CUSTOM","id":"0"},"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %fusion.1136 = f32[128]{0:T(128)S(1)} fusion(%fusion.1133), kind=kLoop, calls=%fused_computation.160.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %add_rsqrt_fusion.13 = f32[128]{0:T(128)S(1)} fusion(%fusion.1136), kind=kLoop, calls=%fused_computation.203.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/rsqrt" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %copy-done.242 = bf16[3,128]{1,0:T(4,128)(2,1)S(1)} copy-done(%copy-start.242) + %fusion.1137 = bf16[1,128]{1,0:T(2,128)(2,1)S(1)} fusion(%copy-done.242, %copy.1644), kind=kLoop, calls=%fused_computation.201.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %all-gather.346 = bf16[1,512]{1,0:T(2,128)(2,1)S(1)} all-gather(%fusion.1137), channel_id=285, replica_groups=mesh['axis_0'=4] {'axis_0'}, dimensions={1}, use_global_device_ids=true, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/...k,k->...k/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"barrier_config":{"barrier_type":"CUSTOM","id":"0"},"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %reduce.885 = bf16[512]{0:T(512)(128)(2,1)S(1)} reduce(%all-gather.346, %constant.4192..sunk.1), dimensions={0}, to_apply=%all-gather.293.reduce_sub_computation, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/...k,k->...k/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %fusion.1138 = bf16[1,128,128,256]{3,1,2,0:T(8,128)(2,1)S(1)} fusion(%all-gather.345, %add_rsqrt_fusion.13, %fusion.1133, %reduce.885), kind=kOutput, calls=%fused_computation.130.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %get-tuple-element.8936 = s32[128]{0:T(128)} get-tuple-element(%wide.param.5), index=55 + %copy-start.216 = (s32[128]{0:T(128)S(1)}, s32[128]{0:T(128)}, u32[]{:S(2)}) copy-start(%get-tuple-element.8936) + %slice_bitcast_fusion.11 = (bf16[128,128,128]{2,1,0:T(8,128)(2,1)S(1)}, bf16[1,128,128,128]{3,1,2,0:T(8,128)(2,1)S(1)}) fusion(%fusion.1138), kind=kLoop, calls=%fused_computation.132.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/vmap(static_sparsity_splash)/jit(_splash_attention)/splash_mha_fwd_segmented_residuals/splash_mha_fwd_segmented_residuals/squeeze" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %get-tuple-element.8755 = bf16[1,128,128,128]{3,1,2,0:T(8,128)(2,1)S(1)} get-tuple-element(%slice_bitcast_fusion.11), index=1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/vmap(static_sparsity_splash)/jit(_splash_attention)/splash_mha_fwd_segmented_residuals/splash_mha_fwd_segmented_residuals/squeeze" stack_frame_id=0} + %bitcast.2582 = s8[16777216]{0:T(1024)(128)(4,1)} bitcast(%all-gather.340) + %slice-start.21 = ((s8[16777216]{0:T(1024)(128)(4,1)}), s8[4194304]{0:T(1024)(128)(4,1)S(1)}, s32[]{:S(2)}) async-start(%bitcast.2582), calls=%async_computation.48 + %slice-start.22 = ((s8[16777216]{0:T(1024)(128)(4,1)}), s8[4194304]{0:T(1024)(128)(4,1)S(1)}, s32[]{:S(2)}) async-start(%bitcast.2582), calls=%async_computation.49 + %maximum_bitcast_fusion.28 = bf16[128,128,192]{2,1,0:T(8,128)(2,1)S(1)} fusion(%broadcast_in_dim.1649, %get-tuple-element.8755), kind=kLoop, calls=%fused_computation.120.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/vmap(static_sparsity_splash)/jit(_splash_attention)/splash_mha_fwd_segmented_residuals/splash_mha_fwd_segmented_residuals/squeeze" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %get-tuple-element.8756 = bf16[128,128,128]{2,1,0:T(8,128)(2,1)S(1)} get-tuple-element(%slice_bitcast_fusion.11), index=0, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/vmap(static_sparsity_splash)/jit(_splash_attention)/splash_mha_fwd_segmented_residuals/splash_mha_fwd_segmented_residuals/squeeze" stack_frame_id=0} + %copy-done.216 = s32[128]{0:T(128)S(1)} copy-done(%copy-start.216) + %squeeze.946 = s32[128,128]{1,0:T(8,128)S(1)} broadcast(%copy-done.216), dimensions={0}, metadata={op_name="jit(train_step)/shard_map/vmap(static_sparsity_splash)/jit(_splash_attention)/splash_mha_fwd_segmented_residuals/splash_mha_fwd_segmented_residuals/squeeze" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"output_chunk_bound_config":{"output_chunk_bound":["8"]},"aliasing_operands":{"lists":[]}} + %squeeze.947 = s32[8,128]{1,0:T(8,128)S(1)} broadcast(%copy-done.216), dimensions={1}, metadata={op_name="jit(train_step)/shard_map/vmap(static_sparsity_splash)/jit(_splash_attention)/splash_mha_fwd_segmented_residuals/splash_mha_fwd_segmented_residuals/squeeze" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"output_chunk_bound_config":{"output_chunk_bound":["128"]},"aliasing_operands":{"lists":[]}} + %iota.360 = s32[128,128]{1,0:T(8,128)S(1)} iota(), iota_dimension=0, metadata={op_name="jit(train_step)/shard_map/vmap(static_sparsity_splash)/jit(_splash_attention)/broadcast_in_dim" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %slice-start.23 = ((s8[16777216]{0:T(1024)(128)(4,1)}), s8[4194304]{0:T(1024)(128)(4,1)S(1)}, s32[]{:S(2)}) async-start(%bitcast.2582), calls=%async_computation.50 + %slice-start.24 = ((s8[16777216]{0:T(1024)(128)(4,1)}), s8[4194304]{0:T(1024)(128)(4,1)S(1)}, s32[]{:S(2)}) async-start(%bitcast.2582), calls=%async_computation.51 + %splash_mha_fwd_segmented_residuals.6 = (f32[128,128]{1,0:T(8,128)}, f32[128,128]{1,0:T(8,128)}, f32[128,128]{1,0:T(8,128)}, bf16[128,128,128]{2,1,0:T(8,128)(2,1)S(1)}, f32[128,128,128]{2,1,0:T(8,128)}) custom-call(%get-tuple-element.8938, %get-tuple-element.8939, %fusion.1132, %maximum_bitcast_fusion.28, %get-tuple-element.8756, /*index=5*/%squeeze.946, %squeeze.947, %iota.360), custom_call_target="tpu_custom_call", operand_layout_constraints={s8[1,1,1]{2,1,0}, s8[1,1,1]{2,1,0}, bf16[128,128,192]{2,1,0}, bf16[128,128,192]{2,1,0}, bf16[128,128,128]{2,1,0}, s32[128,128]{1,0}, s32[8,128]{1,0}, s32[128,128]{1,0}}, frontend_attributes={kernel_metadata={ +"xprof_metadata":"{\"block_q\": 128, \"block_kv\": 128, \"block_kv_compute\": 128, \"block_q_dkv\": 128, \"block_kv_dkv\": 128, \"block_kv_dkv_compute\": 128, \"block_q_dq\": 128, \"block_kv_dq\": 128, \"use_fused_bwd_kernel\": false, \"q_layout\": 1, \"k_layout\": 1, \"v_layout\": 1}" +}}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/vmap(static_sparsity_splash)/jit(_splash_attention)/splash_mha_fwd_segmented_residuals/splash_mha_fwd_segmented_residuals/pallas_call" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"custom_call_config":{"body":"TUzvUgFNTElSMjMuMC4wZ2l0AAFTCwEDBQcJAQMLAz0NDxETFRcZGx0fISMlJykrLS8xMzU3OTs9P0FDRUcDugQGBDkB9QcXExcXEw8XFxsLCxcLCwsTJw8LFxMTUwsXG6UTFwsTExeFcwsTFxcTCxcPC3MLCwsLCwsLCwsLCwsTCwsnCwsLEwsXFwsLCwsXFwsLMxcXFxcTCw9LCxcLC3MLDwsLCws3GwsbCxsLGwsbC1MbCyMLIwsjCxsLGwsbBRFhgZGNYU0qAioCAfcbGxsbGxsbGxsbDw8TCwsrCwsTCxMLDxMLEwsTCwsTCw8TCxcfCxMLCxMLCxMLUxcLEwsLEwsPEwsPUwsTCxMTCwsfCxMLLxsLIxcLEwsXDxcLEwsXCxMLHxcLEwsTDxcLEwsXHwsTCx8LEwsLEwsLEwsXHwsTCx8vFw8HBVlZCQVdSQE5DwcfIwcPCwsfBycXKysnHyMfIysjIx8bG08vKwLSGh8DAx82Ah1xjgIdtgK6AgMDHzICHXE6AhUZGR1+AoICHU4CUgIDA6oC9gMFSQVLHUICRgIFTQVPBVEdU+4CAwd9JX/yAoGDEQ0ADTcd1gLaAh1T5gIdeQYDIwsFIYAAAAAAAAAAgAAAAAAAAAANNQMDH+ICAwMOA/oDYWZmaW5lX21hcDwoZDAsIGQxLCBkMikgLT4gKGQwLCBkMSwgZDIpPgAdc1oCAwM9YgIFUx1zZgIdd3YCHRYDGgNhZmZpbmVfbWFwPChkMCwgZDEpIC0+IChkMCwgZDEpPgAjCwcxAQAAAAAAAACAAAAAAAAAAIAAAAAAAAAABVUDAz1XAwM9igIdngKiAh15rgIFVwMDHwIDEQsBBVkjCwcxAQAAAAAAAACAAAAAAAAAAMAAAAAAAAAABVsFXQVfBWEFYwVlBWcFaQVrBW0FbwVxHXduAgVzBXUDB30lf8ICgYMFdwV5BXsjAQEBBX0d9gL6AgMDHxIDBX8FgQWDBYUDAx9qAx2eA6IDBYcFiSMLAxEBAAAAAAAAAB2uA7IDHboDvgMdxgPKAx3aA94DAwOnqQWLEQs9AxGtrxWxs7W3ubtXF72//8HDBY0BB/v5+Q0zBY8jCwcxgAAAAAAAAAABAAAAAAAAAAEAAAAAAAAABZERCwkFkwWVBZcFmQEXxcnN0dXb3+Pn6+8DBRvHHVsJWQMFG8sdWwldAwUbzx1HCV8DBRvTHS8JYQMFG9cd2QljIwsFIQgAAAAAAAAAgAAAAAAAAAADBRvdHS8JZQMHSfcb4R0vCWcDB0n3G+UdLwlpAwdJ9xvpHS8JawMFG+0dRwltAwUb8R1HCW8DBRUxF1kjdHB1Lm1lbW9yeV9zcGFjZTx2bWVtPgAjdHB1LnBpcGVsaW5lX21vZGU8c3luY2hyb25vdXM+ACN0cHUuZGltZW5zaW9uX3NlbWFudGljczxhcmJpdHJhcnk+ACN0cHUuZGltZW5zaW9uX3NlbWFudGljczxwYXJhbGxlbD4AI3RwdS5tZW1vcnlfc3BhY2U8c21lbT4AI3RwdS5jb3JlX3R5cGU8dGM+ACN0cHUuZG90X2RpbWVuc2lvbl9udW1iZXJzPFsxXSwgWzFdLCBbMF0sIFswXSwgWzAsIDAsIDEsIDBdLCBbXSwgW10+ACN0cHUuZG90X2RpbWVuc2lvbl9udW1iZXJzPFsxXSwgWzBdLCBbMF0sIFsxXSwgWzAsIDAsIDEsIDFdLCBbXSwgW10+AAMFFTEXXQMFFTEXXwMFFScXYQMFFScXYwMFFScXZQMFFScXZwMFFScXaQMFFScXawMFFTEXbQMFFTEXbxEBAREDAR0+AhkFmwWdLUoCCe4lEzImBwWfBaEdVgINBaMdXgIZBaURCxEdagINBacdcgIZBakdegINBasFrR2GAhkFrxELBR2SAg0FsQMDH5oCEwkQAADgDwWzHaYCGQW1BbcdsgINBbkFux2+Ag0FvSMBCSEBAAAAAQAAAAMAAAAAAAAAHcoCzgIFvx3SAg0FwQXDHd4CDQXFEwkBHeoCGQXHHYUNIwEJIQEAAAABAAAAAgAAAAAAAAAFyR3+AhkFyxEBAgQdCgMZBc0FzyUFCQAAAAAF0R0eAw0F0wMJiwICjSWPJZElAwMqAy4DBdUjAQMJAQAAAB02AzoDBdcdPgMNBdkDAz1GAxELFR1OA1IDBdsdVgMNBd0dXgNiAwXfHWYDDQXhEwmQzMzMPx1yA3YDBeMdegMZBeUdU4IDHYUZHYoDjgMF5x2SAwEF6QMDH5oDJRcJAACA/wXrHaYDDQXtAwWX/gOZmwXvHbYDDQXxBfMdwgMNBfUF9x3OAw0F+QMDH9YDJRcJAAAAAAX7HeIDDQX9AwWXAgSZmwMJiwYCjSWPJZElAwMf8gMRAQUjYXJpdGguZmFzdG1hdGg8bm9uZT4AI2FyaXRoLm92ZXJmbG93PG5vbmU+ACN2ZWN0b3Iua2luZDxtYXhpbXVtZj4AI3ZlY3Rvci5raW5kPGFkZD4AAQICAycFAgQCBAkX/QcFBQUPNwsBAgQBCQFBJwUCBAIEAQcX9QUCBAIECUUnAwIECRf1BwUCBAIGEzcX9QcFAgQCBBM3F/UFAgQCBAFFJwUCBAIEEycHBQIEAgQTJwUCBAIEDRf1BSECBAFFF/UHBQIEAgQJNycHBQIEAgQJJwcFAgQCBhMnBQIEAgYTJwUFAgQBJwUCBAUJBSEBAQEHBxkZGx0lHRUVFRsnAQULAQEBBwcHAQEBBQsBAQEHBwUBAQRuIgUBEQGlBwMBMQsRAasHA1N3IQEBAQEBAQcBBwEZARkBGwEdASUBHQEVARUBFQEbAScBAwN1CQMBEQdBSwMNBQUhGwYRAwEDIwMDDwkDAREHD00DDQUlJx0UDwMpCQMfTQMDKzMDCQ8GIQMFA1MDAwcDAwMDAwcDAwMFBgcDBQcbV1kTBQcjCVUbV1kDAyuTAwkPBiEDBQNdAwMHAwMDAwMHAwMDBQYHAwUHF2FjEwUHIwlfF2FjAwMrMwMJDwYhAwUDZwMDBwMDAwMDBwMDAwUGBwMFBxlrbRMFByMJaRlrbRkADwMBBRkADwMDCwMDAwcGCwMDAwMHBgsDAwMFFQYLAw8JCSstLxcGEQMBAzEDAzkJAwERBz87Aw0FMzUDAwsDAwMHBgsDAwMDBwYLAwMDBRUGCwMPCQc5Oz0XBhEDAQM/GwYRAwEDNwMDDwkDAREHD00DDQVDRR0UDwNHCQOz6gIDA4cJAwEDAy1VAwEjBy01AwEFU1UDAwUDAwMDAwUDAwMFBgUDBQcXWVsJBgUDBQNdAwMFAwMDAwMFAwMDBQYFAwUHGWFjCQYFAwUDZQMDBQMDAwMDBQMDAwMDBQMDAwUGBQMrCQtpa20JBgUDLQNvAwMFAwMDBwYFAwMDVwMDBQMDAwUGBQMrCQ1zdXcJBgUDLQN5AwNDiQMFJQdDIgMDBQdxe30DAy1VAwEjBy01AwEFQYEDAy1VAwEjBy01AwEFU4UnByk1AwEFg4c1AzIDJgMDEQ8GKQMRA4knByk1AxEFjYsDAwUDAwMDAwUDAwMFBgUDEQcVkZMJBgUDEQOVEQdKA0IDAyMFl48DAwUDAwMHBgUDAwNXBQYFAy8HE5udCQYFAy8DnwMDBQMDAwMDBQMDAwUGBQMRBxGjpQkGBQMRA6cPBkEDEQOhEQdBSwMjBamrNwZaAwMjBZmtAwNuA5MDCQ8GfgMDBQOxOQaGAwMFB69/swMDlZYDAxcpB5WqAwMXBbW3CQYhAzEDuQ8GnQMFA7s7B50TAwUFX70rB58TAwUFtb8tB6ETAwUDwQMDo9IDAxcpB6PmAwMXBcPFCQYhAzEDxw8GIQMFA8krB58TAwUFX78tB6ETAwUDzR8HURMDBQXPZyEHKRMDBQXL0QMDBwMDAwMDBwMDAwUGBwMFBxfV1xMFByMJvxfV1wMDBwMDAwMDBwMDAwUGBwMFBxnb3RMFByMJ0xnb3QMDBQMDAwcGBQMDA1cDAwUDAwMFBgUDIQkP4ePlCQYFAx8D5z0GEQMFA+kDA0OJAwUlB0PqAwMFB8Pr7QMDBQMDAwMDBQMDAwUGBQMFBxvx8wkGBQMFA/UfB1ETAwUFz/chBykTAwUF+e8DAwcDAwMDAwcDAwMFBgcDBQcb/f8TBQcjCfsb/f8DA4fuAwMBGQAPAwEFGQAPAwN1CQMBEQdBSwMNBQVJGwYRAwEDSwMDDwkDAREHD00DDQVNTx0UDwNRCQNd0QMDBQMDAwMDBQMDAwUGBQMFBxlTVQkGBQMFA1cDA0+WAgMJDwZPAwUDWy8HTxMDBQVdWQMDBQMDAwMDBQMDAwUGBQMFBxthYwkGBQMFA2UfB1ETAwUFZ18xBhEDHwNpAwMHAwMDAwMHAwMDAwMHAwMDBQYHAyEJHW1vcQkGBwMfA3MJBgcDIQNrEwUHewt3HW1vcTMHxgITAwUDWQMDBQMDAwMDBQMDAwUGBQMFBxd7fQkGBQMFA38hBykTAwUFeYEDAwcDAwMDAwcDAwMDAwcDAwMFBgcDKQkfhYeJCQYHAwUDiwkGBwMpA4MTBQd7C48fhYeJAwMrMwMJDwYhAwUDkQMDBwMDAwMDBwMDAwUGBwMFBxeVlxMFByMJkxeVlwMDKzMDCQ8GIQMFA5sDAwcDAwMDAwcDAwMFBgcDBQcZn6ETBQcjCZ0Zn6EDAyszAwkPBiEDBQOlAwMHAwMDAwMHAwMDBQYHAwUHG6mrEwUHIwmnG6mrGQAPAwEFGQAPDQABCxEB8wcDDw8LAQEBAQEBBwEHAQMDAQkDAQMDAQkDAQ0EAQcBAwsLEQEKAgcDJz8LAQEBAQEBBwEHAQMDCwMDAwcGCwMDAwMHBgsDAwMFFQYLAw8JCQsNDxcGEQMBAxEDAzkJAwERBz87Aw0FExUDAwsDAwMHBgsDAwMDBwYLAwMDBRUGCwMPCQcZGx0XBhEDAQMfAwMBCQMBAwMBCQMBDQQBBwEhIwsRAQ4CBwMnPwsBAQEBAQEHAQcBAwMLAwMDBwYLAwMDAwcGCwMDAwUVBgsDDwkJCw0PFwYRAwEDEQMDOQkDAREHPzsDDQUTFQMDCwMDAwcGCwMDAwMHBgsDAwMFFQYLAw8JBxkbHRcGEQMBAx8DAwEJAwEDAwEJAwENBAEHASEjCxEBEgIHAw8PCwEBAQEBAQcBBwEDAwEJAwEDAwEJAwENBAEFAwsLEQEWAgcDJz8LAQEBAQEBBwEHAQMDCwMDAwcGCwMDAwMHBgsDAwMFFQYLAw8JCQsNDxcGEQMBAxEDAzkJAwERBz87Aw0FExUDAwsDAwMHBgsDAwMDBwYLAwMDBRUGCwMPCQcZGx0XBhEDAQMfAwMBCQMBAwMBCQMBDQQBBSMhCxEBGgIHAw8PCwEBAQEBAQcBBwEDAwEJAwEDAwEJAwENBAEFAwsLEQEeAgcDERMLAQEBAQEBBwEHAQMDAQkDAQMDAQkDAQMDAQkDAQ0EAQULDQsRASICBwMREwsBAQEBAQEHAQcBAwMBCQMBAwMBCQMBAwMBCQMBDQQBBQsNCxEBJgIHAxETCwEBAQEBAQcBBwEDAwEJAwEDAwEJAwEDAwEJAwENBAEFCw0LEQEqAgcDDw8LAQEBAQEBBwEHAQMDAQkDAQMDAQkDAQ0EAQcBAwsLEQEuAgcDDw8LAQEBAQEBBwEHAQMDAQkDAQMDAQkDAQ0EAQcBAwsGAwEFAQCuGv8tGRELEQsRCy0ZJRUJCxELDQkVDRcxGx0JCw0jEQsRCxUNERMJCxELDQ0HDQdVLRICJQkdHUcjISMpLR8LHScdJUURKQkLCQkLGxkZGRkZGRkZGRklHRUNHSUTHRcfGxcTExsXExcvFxcXFxcPGRUZGSMXIxkVJSMZHw8PDQkdEWJ1aWx0aW4Ac3RhYmxlX21vc2FpYwB0cHUAYXJpdGgAdmVjdG9yAG1vZHVsZQBhcml0aC5jb25zdGFudAB2ZWN0b3IubG9hZABhcml0aC5pbmRleF9jYXN0AHZlY3Rvci5zaGFwZV9jYXN0AGZ1bmMuZnVuYwBmdW5jLnJldHVybgB2ZWN0b3IuYnJvYWRjYXN0AGFyaXRoLmNtcGkAdHB1LnZlY3Rvcl9zdG9yZQBtZW1yZWYubG9hZABhcml0aC5leHRzaQBzY2YueWllbGQAYXJpdGguZXh0dWkAc2NmLmlmAGFyaXRoLm11bGYAYXJpdGguYWRkZgBhcml0aC5tdWxpAHRwdS5tYXRtdWwAYXJpdGguYWRkaQB2ZWN0b3IubXVsdGlfcmVkdWN0aW9uAGFyaXRoLnN1YmYAbWF0aC5leHAAYXJpdGguZGl2ZgBhcml0aC50cnVuY2YAbWF0aC5sb2cAdHB1LmlvdGEAYXJpdGguYW5kaQBhcml0aC5zZWxlY3QAYXJpdGgubWF4aW11bWYAYXJpdGguZXh0ZgBmdW5jdGlvbl90eXBlAHN5bV9uYW1lAHRyYW5zZm9ybV9pbmRpY2VzAHdpbmRvd19ib3VuZHMAdmFsdWUAcHJlZGljYXRlAHBpcGVsaW5lX21vZGUAYnJvYWRjYXN0X2luX2RpbToAdHJhbnNmb3JtXzAAdHJhbnNmb3JtXzEAdHJhbnNmb3JtXzIAdHJhbnNmb3JtXzMAdHJhbnNmb3JtXzQAdHJhbnNmb3JtXzUAdHJhbnNmb3JtXzYAdHJhbnNmb3JtXzcAdHJhbnNmb3JtXzgAdHJhbnNmb3JtXzkAdHJhbnNmb3JtXzEwAGdldDoAZ3Q6AGVxOgBtdWw6AGFkZABvcGVyYW5kU2VnbWVudFNpemVzAHN0cmlkZXMAYnJvYWRjYXN0X2luX2RpbS9icm9hZGNhc3RfaW5fZGltAGRpbWVuc2lvbl9udW1iZXJzAHRyYW5zcG9zZV9saHMAdHJhbnNwb3NlX2xoc19oaW50AHRyYW5zcG9zZV9yaHMAa2luZAByZWR1Y3Rpb25fZGltcwBzdGFibGVfbW9zYWljLnZlcnNpb24AZGltZW5zaW9uX3NlbWFudGljcwBpdGVyYXRpb25fYm91bmRzAHNjYWxhcl9wcmVmZXRjaABzY3JhdGNoX29wZXJhbmRzAHNwbGFzaF9taGFfZndkX3NlZ21lbnRlZF9yZXNpZHVhbHMAdHB1LmNvcmVfdHlwZQB3aW5kb3dfcGFyYW1zAGdldABfc3BsYXNoX2F0dGVudGlvbgAvX193L21heHRleHQvbWF4dGV4dC8udmVudi9saWIvcHl0aG9uMy4xMi9zaXRlLXBhY2thZ2VzL2pheC9leHBlcmltZW50YWwvcGFsbGFzL29wcy90cHUvc3BsYXNoX2F0dGVudGlvbi9zcGxhc2hfYXR0ZW50aW9uX2tlcm5lbC5weQBjb252ZXJ0X2VsZW1lbnRfdHlwZToAY29udmVydF9lbGVtZW50X3R5cGUvY29udmVydF9lbGVtZW50X3R5cGUAZ3QAZ3QvZ3QAZXEAZXEvZXEAY29uZDoAY29uZABnZXQvZ2V0AGRpdjoAZGl2AGZhc3RtYXRoAG11bC9tdWwAc3dhcDoAc3dhcC9zd2FwAGxvZzoAbG9nL2xvZwBhZGQ6AGFkZC9hZGQAYnJvYWRjYXN0X2luX2RpbQBzY2FuOgBzY2FuAG11bABvdmVyZmxvd0ZsYWdzAGRvdF9nZW5lcmFsOgBkb3RfZ2VuZXJhbC9kb3RfZ2VuZXJhbABkaW1lbnNpb25zAGlvdGE6AGlvdGEvaW90YQBnZToAZ2UvZ2UAYW5kOgBhbmQvYW5kAGppdDoAaml0AHNlbGVjdF9uOgBzZWxlY3Rfbi9zZWxlY3RfbgByZWR1Y2VfbWF4OgByZWR1Y2VfbWF4L3JlZHVjZV9tYXgAbWF4OgBtYXgvbWF4AHN1YjoAc3ViL3N1YgBleHA6AGV4cC9leHAAcmVkdWNlX3N1bToAcmVkdWNlX3N1bS9yZWR1Y2Vfc3VtAA==","needs_layout_passes":true,"allow_input_fusion":[],"serialization_format":"1","output_memory_colors":[],"output_memory_space_colors":[],"input_memory_space_colors":[]},"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %get-tuple-element.8757 = bf16[128,128,128]{2,1,0:T(8,128)(2,1)S(1)} get-tuple-element(%splash_mha_fwd_segmented_residuals.6), index=3, frontend_attributes={kernel_metadata={ +"xprof_metadata":"{\"block_q\": 128, \"block_kv\": 128, \"block_kv_compute\": 128, \"block_q_dkv\": 128, \"block_kv_dkv\": 128, \"block_kv_dkv_compute\": 128, \"block_q_dq\": 128, \"block_kv_dq\": 128, \"use_fused_bwd_kernel\": false, \"q_layout\": 1, \"k_layout\": 1, \"v_layout\": 1}" +}}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/vmap(static_sparsity_splash)/jit(_splash_attention)/splash_mha_fwd_segmented_residuals/splash_mha_fwd_segmented_residuals/pallas_call" stack_frame_id=0} + %slice-done.21 = s8[4194304]{0:T(1024)(128)(4,1)S(1)} async-done(%slice-start.21) + %slice-done.22 = s8[4194304]{0:T(1024)(128)(4,1)S(1)} async-done(%slice-start.22) + %slice-done.23 = s8[4194304]{0:T(1024)(128)(4,1)S(1)} async-done(%slice-start.23) + %slice-done.24 = s8[4194304]{0:T(1024)(128)(4,1)S(1)} async-done(%slice-start.24) + %custom-call.91 = bf16[1,128,128,512]{3,2,1,0:T(8,128)(2,1)S(1)} custom-call(%slice-done.21, %slice-done.22, %slice-done.23, %slice-done.24), custom_call_target="ConcatBitcast", backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %bitcast.2583 = bf16[1,128,128,512]{3,2,1,0:T(8,128)(2,1)S(1)} bitcast(%custom-call.91) + %fusion.1139 = (f32[128]{0:T(128)S(1)}, bf16[1,128,512]{2,1,0:T(8,128)(2,1)S(1)}) fusion(%get-tuple-element.8822, %bitcast.2583, %get-tuple-element.8757), kind=kOutput, calls=%fused_computation.162.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %get-tuple-element.8758 = bf16[1,128,512]{2,1,0:T(8,128)(2,1)S(1)} get-tuple-element(%fusion.1139), index=1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/reduce_sum" stack_frame_id=0} + %get-tuple-element.8915 = bf16[3,18432,128]{2,1,0:T(8,128)(2,1)} get-tuple-element(%wide.param.5), index=34 + %dynamic-slice_convert_fusion.45 = bf16[1,18432,128]{2,1,0:T(8,128)(2,1)S(1)} fusion(%get-tuple-element.8915, %copy.1644), kind=kLoop, calls=%fused_computation.115.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %all-gather.347 = bf16[1,18432,512]{2,1,0:T(8,128)(2,1)S(1)} all-gather(%dynamic-slice_convert_fusion.45), channel_id=291, replica_groups=mesh['axis_0'=1,'axis_1'=4] {'axis_1'}, dimensions={2}, use_global_device_ids=true, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"barrier_config":{"barrier_type":"CUSTOM","id":"0"},"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %get-tuple-element.8914 = bf16[3,128,18432]{2,1,0:T(8,128)(2,1)} get-tuple-element(%wide.param.5), index=33 + %dynamic-slice_convert_fusion.46 = bf16[1,128,18432]{2,1,0:T(8,128)(2,1)S(1)} fusion(%get-tuple-element.8914, %copy.1644), kind=kLoop, calls=%fused_computation.116.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %all-gather.348 = bf16[1,512,18432]{2,1,0:T(8,128)(2,1)S(1)} all-gather(%dynamic-slice_convert_fusion.46), channel_id=290, replica_groups=mesh['axis_0'=4,'axis_1'=1] {'axis_0'}, dimensions={1}, use_global_device_ids=true, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"barrier_config":{"barrier_type":"CUSTOM","id":"0"},"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %get-tuple-element.8913 = bf16[3,128,18432]{2,1,0:T(8,128)(2,1)} get-tuple-element(%wide.param.5), index=32 + %dynamic-slice_convert_fusion.47 = bf16[1,128,18432]{2,1,0:T(8,128)(2,1)S(1)} fusion(%get-tuple-element.8913, %copy.1644), kind=kLoop, calls=%fused_computation.117.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %all-gather.349 = bf16[1,512,18432]{2,1,0:T(8,128)(2,1)S(1)} all-gather(%dynamic-slice_convert_fusion.47), channel_id=289, replica_groups=mesh['axis_0'=4,'axis_1'=1] {'axis_0'}, dimensions={1}, use_global_device_ids=true, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"barrier_config":{"barrier_type":"CUSTOM","id":"0"},"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %get-tuple-element.8759 = f32[128]{0:T(128)S(1)} get-tuple-element(%fusion.1139), index=0, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/reduce_sum" stack_frame_id=0} + %add_rsqrt_fusion.14 = f32[128]{0:T(128)S(1)} fusion(%get-tuple-element.8759), kind=kLoop, calls=%fused_computation.202.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/rsqrt" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %get-tuple-element.8760 = bf16[512]{0:T(512)(128)(2,1)S(1)} get-tuple-element(%fusion.1128), index=0, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %get-tuple-element.8916 = u32[3]{0:T(128)} get-tuple-element(%wide.param.5), index=35, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_update_index_in_dim)/dynamic_update_slice" stack_frame_id=0} + %copy-start.304 = (u32[3]{0:T(128)S(1)}, u32[3]{0:T(128)}, u32[]{:S(2)}) copy-start(%get-tuple-element.8916) + %fusion.1140 = bf16[1,128,18432]{2,1,0:T(8,128)(2,1)S(1)} fusion(%all-gather.349, %get-tuple-element.8758, %add_rsqrt_fusion.14, %get-tuple-element.8760), kind=kOutput, calls=%fused_computation.111.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %get-tuple-element.8917 = u32[3,4]{1,0:T(4,128)} get-tuple-element(%wide.param.5), index=36, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_update_index_in_dim)/dynamic_update_slice" stack_frame_id=0} + %copy-start.249 = (u32[3,4]{1,0:T(4,128)S(1)}, u32[3,4]{1,0:T(4,128)}, u32[]{:S(2)}) copy-start(%get-tuple-element.8917) + %get-tuple-element.8918 = u32[3]{0:T(128)} get-tuple-element(%wide.param.5), index=37, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_update_index_in_dim)/dynamic_update_slice" stack_frame_id=0} + %copy-start.305 = (u32[3]{0:T(128)S(1)}, u32[3]{0:T(128)}, u32[]{:S(2)}) copy-start(%get-tuple-element.8918) + %get-tuple-element.8919 = u32[3,4]{1,0:T(4,128)} get-tuple-element(%wide.param.5), index=38, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_update_index_in_dim)/dynamic_update_slice" stack_frame_id=0} + %copy-start.250 = (u32[3,4]{1,0:T(4,128)S(1)}, u32[3,4]{1,0:T(4,128)}, u32[]{:S(2)}) copy-start(%get-tuple-element.8919) + %get-tuple-element.8920 = u32[3]{0:T(128)} get-tuple-element(%wide.param.5), index=39, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_update_index_in_dim)/dynamic_update_slice" stack_frame_id=0} + %copy-start.306 = (u32[3]{0:T(128)S(1)}, u32[3]{0:T(128)}, u32[]{:S(2)}) copy-start(%get-tuple-element.8920) + %get-tuple-element.8921 = u32[3,4]{1,0:T(4,128)} get-tuple-element(%wide.param.5), index=40, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_update_index_in_dim)/dynamic_update_slice" stack_frame_id=0} + %copy-start.251 = (u32[3,4]{1,0:T(4,128)S(1)}, u32[3,4]{1,0:T(4,128)}, u32[]{:S(2)}) copy-start(%get-tuple-element.8921) + %get-tuple-element.8922 = u32[3]{0:T(128)} get-tuple-element(%wide.param.5), index=41, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_update_index_in_dim)/dynamic_update_slice" stack_frame_id=0} + %copy-start.307 = (u32[3]{0:T(128)S(1)}, u32[3]{0:T(128)}, u32[]{:S(2)}) copy-start(%get-tuple-element.8922) + %get-tuple-element.8923 = u32[3,4]{1,0:T(4,128)} get-tuple-element(%wide.param.5), index=42, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_update_index_in_dim)/dynamic_update_slice" stack_frame_id=0} + %copy-start.252 = (u32[3,4]{1,0:T(4,128)S(1)}, u32[3,4]{1,0:T(4,128)}, u32[]{:S(2)}) copy-start(%get-tuple-element.8923) + %get-tuple-element.8924 = u32[3]{0:T(128)} get-tuple-element(%wide.param.5), index=43, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_update_index_in_dim)/dynamic_update_slice" stack_frame_id=0} + %copy-start.308 = (u32[3]{0:T(128)S(1)}, u32[3]{0:T(128)}, u32[]{:S(2)}) copy-start(%get-tuple-element.8924) + %get-tuple-element.8925 = u32[3,4]{1,0:T(4,128)} get-tuple-element(%wide.param.5), index=44, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_update_index_in_dim)/dynamic_update_slice" stack_frame_id=0} + %copy-start.253 = (u32[3,4]{1,0:T(4,128)S(1)}, u32[3,4]{1,0:T(4,128)}, u32[]{:S(2)}) copy-start(%get-tuple-element.8925) + %get-tuple-element.8926 = u32[3]{0:T(128)} get-tuple-element(%wide.param.5), index=45, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_update_index_in_dim)/dynamic_update_slice" stack_frame_id=0} + %copy-start.309 = (u32[3]{0:T(128)S(1)}, u32[3]{0:T(128)}, u32[]{:S(2)}) copy-start(%get-tuple-element.8926) + %get-tuple-element.8927 = u32[3,4]{1,0:T(4,128)} get-tuple-element(%wide.param.5), index=46, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_update_index_in_dim)/dynamic_update_slice" stack_frame_id=0} + %copy-start.254 = (u32[3,4]{1,0:T(4,128)S(1)}, u32[3,4]{1,0:T(4,128)}, u32[]{:S(2)}) copy-start(%get-tuple-element.8927) + %get-tuple-element.8928 = u32[3]{0:T(128)} get-tuple-element(%wide.param.5), index=47, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_update_index_in_dim)/dynamic_update_slice" stack_frame_id=0} + %copy-start.310 = (u32[3]{0:T(128)S(1)}, u32[3]{0:T(128)}, u32[]{:S(2)}) copy-start(%get-tuple-element.8928) + %get-tuple-element.8929 = u32[3,4]{1,0:T(4,128)} get-tuple-element(%wide.param.5), index=48, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_update_index_in_dim)/dynamic_update_slice" stack_frame_id=0} + %copy-start.255 = (u32[3,4]{1,0:T(4,128)S(1)}, u32[3,4]{1,0:T(4,128)}, u32[]{:S(2)}) copy-start(%get-tuple-element.8929) + %get-tuple-element.8930 = u32[3]{0:T(128)} get-tuple-element(%wide.param.5), index=49, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_update_index_in_dim)/dynamic_update_slice" stack_frame_id=0} + %copy-start.311 = (u32[3]{0:T(128)S(1)}, u32[3]{0:T(128)}, u32[]{:S(2)}) copy-start(%get-tuple-element.8930) + %get-tuple-element.8931 = u32[3,4]{1,0:T(4,128)} get-tuple-element(%wide.param.5), index=50, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_update_index_in_dim)/dynamic_update_slice" stack_frame_id=0} + %copy-start.256 = (u32[3,4]{1,0:T(4,128)S(1)}, u32[3,4]{1,0:T(4,128)}, u32[]{:S(2)}) copy-start(%get-tuple-element.8931) + %get-tuple-element.8932 = u32[3]{0:T(128)} get-tuple-element(%wide.param.5), index=51, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_update_index_in_dim)/dynamic_update_slice" stack_frame_id=0} + %copy-start.312 = (u32[3]{0:T(128)S(1)}, u32[3]{0:T(128)}, u32[]{:S(2)}) copy-start(%get-tuple-element.8932) + %get-tuple-element.8933 = u32[3,4]{1,0:T(4,128)} get-tuple-element(%wide.param.5), index=52, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_update_index_in_dim)/dynamic_update_slice" stack_frame_id=0} + %copy-start.257 = (u32[3,4]{1,0:T(4,128)S(1)}, u32[3,4]{1,0:T(4,128)}, u32[]{:S(2)}) copy-start(%get-tuple-element.8933) + %fusion.1141 = bf16[1,128,512]{2,1,0:T(8,128)(2,1)S(1)} fusion(%get-tuple-element.8758, %all-gather.347, %all-gather.348, %fusion.1140, %add_rsqrt_fusion.14, /*index=5*/%get-tuple-element.8760), kind=kOutput, calls=%fused_computation.146.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %copy-done.304 = u32[3]{0:T(128)S(1)} copy-done(%copy-start.304) + %copy.1667 = u32[3]{0:T(128)} copy(%copy-done.304), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %copy-done.249 = u32[3,4]{1,0:T(4,128)S(1)} copy-done(%copy-start.249) + %copy.1668 = u32[3,4]{1,0:T(4,128)} copy(%copy-done.249), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %copy-done.305 = u32[3]{0:T(128)S(1)} copy-done(%copy-start.305) + %copy.1669 = u32[3]{0:T(128)} copy(%copy-done.305), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %copy-done.250 = u32[3,4]{1,0:T(4,128)S(1)} copy-done(%copy-start.250) + %copy.1670 = u32[3,4]{1,0:T(4,128)} copy(%copy-done.250), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %copy-done.306 = u32[3]{0:T(128)S(1)} copy-done(%copy-start.306) + %copy.1671 = u32[3]{0:T(128)} copy(%copy-done.306), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %copy-done.251 = u32[3,4]{1,0:T(4,128)S(1)} copy-done(%copy-start.251) + %copy.1672 = u32[3,4]{1,0:T(4,128)} copy(%copy-done.251), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %copy-done.307 = u32[3]{0:T(128)S(1)} copy-done(%copy-start.307) + %copy.1673 = u32[3]{0:T(128)} copy(%copy-done.307), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %copy-done.252 = u32[3,4]{1,0:T(4,128)S(1)} copy-done(%copy-start.252) + %copy.1674 = u32[3,4]{1,0:T(4,128)} copy(%copy-done.252), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %copy-done.308 = u32[3]{0:T(128)S(1)} copy-done(%copy-start.308) + %copy.1675 = u32[3]{0:T(128)} copy(%copy-done.308), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %copy-done.253 = u32[3,4]{1,0:T(4,128)S(1)} copy-done(%copy-start.253) + %copy.1676 = u32[3,4]{1,0:T(4,128)} copy(%copy-done.253), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %copy-done.309 = u32[3]{0:T(128)S(1)} copy-done(%copy-start.309) + %copy.1677 = u32[3]{0:T(128)} copy(%copy-done.309), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %copy-done.254 = u32[3,4]{1,0:T(4,128)S(1)} copy-done(%copy-start.254) + %copy.1678 = u32[3,4]{1,0:T(4,128)} copy(%copy-done.254), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %copy-done.310 = u32[3]{0:T(128)S(1)} copy-done(%copy-start.310) + %copy.1679 = u32[3]{0:T(128)} copy(%copy-done.310), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %copy-done.255 = u32[3,4]{1,0:T(4,128)S(1)} copy-done(%copy-start.255) + %copy.1680 = u32[3,4]{1,0:T(4,128)} copy(%copy-done.255), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %copy-done.311 = u32[3]{0:T(128)S(1)} copy-done(%copy-start.311) + %copy.1681 = u32[3]{0:T(128)} copy(%copy-done.311), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %copy-done.256 = u32[3,4]{1,0:T(4,128)S(1)} copy-done(%copy-start.256) + %copy.1682 = u32[3,4]{1,0:T(4,128)} copy(%copy-done.256), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %copy-done.312 = u32[3]{0:T(128)S(1)} copy-done(%copy-start.312) + %copy.1683 = u32[3]{0:T(128)} copy(%copy-done.312), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %copy-done.257 = u32[3,4]{1,0:T(4,128)S(1)} copy-done(%copy-start.257) + %copy.1684 = u32[3,4]{1,0:T(4,128)} copy(%copy-done.257), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %get-tuple-element.8937 = s32[]{:T(128)} get-tuple-element(%wide.param.5), index=56 + %get-tuple-element.8940 = bf16[]{:T(256)} get-tuple-element(%wide.param.5), index=59, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %bitcast.2667 = f32[1,128,1,32]{3,1,2,0:T(8,128)S(1)} bitcast(%get-tuple-element.8934) + %bitcast.2670 = f32[1,128,1,32]{3,1,2,0:T(8,128)S(1)} bitcast(%get-tuple-element.8935) + ROOT %tuple.873 = (s32[]{:T(128)}, bf16[1,128,512]{2,1,0:T(8,128)(2,1)S(1)}, u32[3]{0:T(128)}, u32[3,4]{1,0:T(4,128)}, u32[3]{0:T(128)}, /*index=5*/u32[3,4]{1,0:T(4,128)}, u32[3]{0:T(128)}, u32[3,4]{1,0:T(4,128)}, u32[3]{0:T(128)}, u32[3,4]{1,0:T(4,128)}, /*index=10*/u32[3]{0:T(128)}, u32[3,4]{1,0:T(4,128)}, u32[3]{0:T(128)}, u32[3,4]{1,0:T(4,128)}, u32[3]{0:T(128)}, /*index=15*/u32[3,4]{1,0:T(4,128)}, u32[3]{0:T(128)}, u32[3,4]{1,0:T(4,128)}, u32[3]{0:T(128)}, u32[3,4]{1,0:T(4,128)}, /*index=20*/bf16[3,1,128,512]{3,2,1,0:T(8,128)(2,1)}, bf16[3,512]{1,0:T(4,128)(2,1)}, bf16[3,384]{1,0:T(4,128)(2,1)}, bf16[3,512]{1,0:T(4,128)(2,1)}, bf16[3,128,1536]{2,1,0:T(8,128)(2,1)}, /*index=25*/bf16[3,384,128,192]{2,1,3,0:T(8,128)(2,1)}, bf16[3,128]{1,0:T(4,128)(2,1)}, bf16[3,128,576]{2,1,0:T(8,128)(2,1)S(1)}, bf16[3,128,128,256]{3,1,2,0:T(8,128)(2,1)}, f32[1,128,1,32]{3,1,2,0:T(8,128)S(1)}, /*index=30*/f32[1,128,1,32]{3,1,2,0:T(8,128)S(1)}, bf16[3,128,128,128]{3,2,1,0:T(8,128)(2,1)}, bf16[3,128,18432]{2,1,0:T(8,128)(2,1)}, bf16[3,128,18432]{2,1,0:T(8,128)(2,1)}, bf16[3,18432,128]{2,1,0:T(8,128)(2,1)}, /*index=35*/u32[3]{0:T(128)}, u32[3,4]{1,0:T(4,128)}, u32[3]{0:T(128)}, u32[3,4]{1,0:T(4,128)}, u32[3]{0:T(128)}, /*index=40*/u32[3,4]{1,0:T(4,128)}, u32[3]{0:T(128)}, u32[3,4]{1,0:T(4,128)}, u32[3]{0:T(128)}, u32[3,4]{1,0:T(4,128)}, /*index=45*/u32[3]{0:T(128)}, u32[3,4]{1,0:T(4,128)}, u32[3]{0:T(128)}, u32[3,4]{1,0:T(4,128)}, u32[3]{0:T(128)}, /*index=50*/u32[3,4]{1,0:T(4,128)}, u32[3]{0:T(128)}, u32[3,4]{1,0:T(4,128)}, f32[128,32]{1,0:T(8,128)S(1)}, f32[128,32]{1,0:T(8,128)S(1)}, /*index=55*/s32[128]{0:T(128)}, s32[]{:T(128)}, s8[1,1,1]{2,1,0:T(4,128)(4,1)}, s8[1,1,1]{2,1,0:T(4,128)(4,1)}, bf16[]{:T(256)}) tuple(%add.3975, %fusion.1141, %copy.1667, %copy.1668, %copy.1669, /*index=5*/%copy.1670, %copy.1671, %copy.1672, %copy.1673, %copy.1674, /*index=10*/%copy.1675, %copy.1676, %copy.1677, %copy.1678, %copy.1679, /*index=15*/%copy.1680, %copy.1681, %copy.1682, %copy.1683, %copy.1684, /*index=20*/%bitcast_dynamic-update-slice_fusion.17, %get-tuple-element.8902, %get-tuple-element.8903, %get-tuple-element.8904, %get-tuple-element.8905, /*index=25*/%get-tuple-element.8906, %get-tuple-element.8907, %get-tuple-element.8908, %get-tuple-element.8909, %bitcast.2667, /*index=30*/%bitcast.2670, %get-tuple-element.8912, %get-tuple-element.8913, %get-tuple-element.8914, %get-tuple-element.8915, /*index=35*/%get-tuple-element.8916, %get-tuple-element.8917, %get-tuple-element.8918, %get-tuple-element.8919, %get-tuple-element.8920, /*index=40*/%get-tuple-element.8921, %get-tuple-element.8922, %get-tuple-element.8923, %get-tuple-element.8924, %get-tuple-element.8925, /*index=45*/%get-tuple-element.8926, %get-tuple-element.8927, %get-tuple-element.8928, %get-tuple-element.8929, %get-tuple-element.8930, /*index=50*/%get-tuple-element.8931, %get-tuple-element.8932, %get-tuple-element.8933, %get-tuple-element.8934, %get-tuple-element.8935, /*index=55*/%get-tuple-element.8936, %get-tuple-element.8937, %get-tuple-element.8938, %get-tuple-element.8939, %get-tuple-element.8940) +} + +%wide.region_6.9_spmd.clone.clone.clone (wide.param.765: (s32[], bf16[1,128,512], u32[3], u32[3,4], u32[3], /*index=5*/u32[3,4], u32[3], u32[3,4], u32[3], u32[3,4], /*index=10*/u32[3], u32[3,4], u32[3], u32[3,4], u32[3], /*index=15*/u32[3,4], u32[3], u32[3,4], u32[3], u32[3,4], /*index=20*/bf16[3,1,128,512], bf16[3,512], bf16[3,384], bf16[3,512], bf16[3,128,1536], /*index=25*/bf16[3,384,128,192], bf16[3,128], bf16[3,128,576], bf16[3,128,128,256], f32[1,128,1,32], /*index=30*/f32[1,128,1,32], bf16[3,128,128,128], bf16[3,128,18432], bf16[3,128,18432], bf16[3,18432,128], /*index=35*/u32[3], u32[3,4], u32[3], u32[3,4], u32[3], /*index=40*/u32[3,4], u32[3], u32[3,4], u32[3], u32[3,4], /*index=45*/u32[3], u32[3,4], u32[3], u32[3,4], u32[3], /*index=50*/u32[3,4], u32[3], u32[3,4], f32[128,32], f32[128,32], /*index=55*/s32[128], s32[], s8[1,1,1], s8[1,1,1], bf16[])) -> pred[] { + %constant.1499.clone.4 = s32[]{:T(128)} constant(3) + %wide.param.765 = (s32[]{:T(128)}, bf16[1,128,512]{2,1,0:T(8,128)(2,1)S(1)}, u32[3]{0:T(128)}, u32[3,4]{1,0:T(4,128)}, u32[3]{0:T(128)}, /*index=5*/u32[3,4]{1,0:T(4,128)}, u32[3]{0:T(128)}, u32[3,4]{1,0:T(4,128)}, u32[3]{0:T(128)}, u32[3,4]{1,0:T(4,128)}, /*index=10*/u32[3]{0:T(128)}, u32[3,4]{1,0:T(4,128)}, u32[3]{0:T(128)}, u32[3,4]{1,0:T(4,128)}, u32[3]{0:T(128)}, /*index=15*/u32[3,4]{1,0:T(4,128)}, u32[3]{0:T(128)}, u32[3,4]{1,0:T(4,128)}, u32[3]{0:T(128)}, u32[3,4]{1,0:T(4,128)}, /*index=20*/bf16[3,1,128,512]{3,2,1,0:T(8,128)(2,1)}, bf16[3,512]{1,0:T(4,128)(2,1)}, bf16[3,384]{1,0:T(4,128)(2,1)}, bf16[3,512]{1,0:T(4,128)(2,1)}, bf16[3,128,1536]{2,1,0:T(8,128)(2,1)}, /*index=25*/bf16[3,384,128,192]{2,1,3,0:T(8,128)(2,1)}, bf16[3,128]{1,0:T(4,128)(2,1)}, bf16[3,128,576]{2,1,0:T(8,128)(2,1)S(1)}, bf16[3,128,128,256]{3,1,2,0:T(8,128)(2,1)}, f32[1,128,1,32]{3,1,2,0:T(8,128)S(1)}, /*index=30*/f32[1,128,1,32]{3,1,2,0:T(8,128)S(1)}, bf16[3,128,128,128]{3,2,1,0:T(8,128)(2,1)}, bf16[3,128,18432]{2,1,0:T(8,128)(2,1)}, bf16[3,128,18432]{2,1,0:T(8,128)(2,1)}, bf16[3,18432,128]{2,1,0:T(8,128)(2,1)}, /*index=35*/u32[3]{0:T(128)}, u32[3,4]{1,0:T(4,128)}, u32[3]{0:T(128)}, u32[3,4]{1,0:T(4,128)}, u32[3]{0:T(128)}, /*index=40*/u32[3,4]{1,0:T(4,128)}, u32[3]{0:T(128)}, u32[3,4]{1,0:T(4,128)}, u32[3]{0:T(128)}, u32[3,4]{1,0:T(4,128)}, /*index=45*/u32[3]{0:T(128)}, u32[3,4]{1,0:T(4,128)}, u32[3]{0:T(128)}, u32[3,4]{1,0:T(4,128)}, u32[3]{0:T(128)}, /*index=50*/u32[3,4]{1,0:T(4,128)}, u32[3]{0:T(128)}, u32[3,4]{1,0:T(4,128)}, f32[128,32]{1,0:T(8,128)S(1)}, f32[128,32]{1,0:T(8,128)S(1)}, /*index=55*/s32[128]{0:T(128)}, s32[]{:T(128)}, s8[1,1,1]{2,1,0:T(4,128)(4,1)}, s8[1,1,1]{2,1,0:T(4,128)(4,1)}, bf16[]{:T(256)}) parameter(0) + %get-tuple-element.8361 = s32[]{:T(128)} get-tuple-element(%wide.param.765), index=0 + ROOT %lt.873 = pred[]{:T(512)} compare(%get-tuple-element.8361, %constant.1499.clone.4), direction=LT, metadata={op_name="jit(train_step)/jvp()/while/cond/lt" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%fused_computation.677 (param_0.1957: f32[128,1,128,128]) -> bf16[128,1,128,128] { + %param_0.1957 = f32[128,1,128,128]{3,2,1,0:T(8,128)} parameter(0) + %bitcast.1383 = f32[128,1,128,128]{3,2,0,1:T(8,128)} bitcast(%param_0.1957), sharding={devices=[1,1,1,4]<=[4]}, metadata={op_name="state[\'model\'][\'decoder\'][\'moe_layers\'][\'self_attention\'][\'out\'][\'kernel\'].value"} + ROOT %convert_element_type.2569 = bf16[128,1,128,128]{3,2,0,1:T(8,128)(2,1)S(1)} convert(%bitcast.1383), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} +} + +%region_8.13 (reduce_sum.30: f32[], reduce_sum.34: f32[]) -> f32[] { + %reduce_sum.30 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/reduce_sum"} + %reduce_sum.34 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/reduce_sum"} + ROOT %reduce_sum.39 = f32[]{:T(128)} add(%reduce_sum.30, %reduce_sum.34), metadata={op_name="reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%fused_computation.895 (param_0.4494: bf16[1,128,512]) -> f32[128] { + %param_0.4494 = bf16[1,128,512]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.2707 = f32[1,128,512]{2,1,0:T(8,128)} convert(%param_0.4494), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %square.628 = f32[1,128,512]{2,1,0:T(8,128)} multiply(%convert_element_type.2707, %convert_element_type.2707), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/square" stack_frame_id=0} + %constant.1730.clone.94 = f32[]{:T(128)} constant(0) + ROOT %reduce.765 = f32[128]{0:T(128)S(1)} reduce(%square.628, %constant.1730.clone.94), dimensions={0,2}, to_apply=%region_8.13, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/reduce_sum" stack_frame_id=0} +} + +%fused_computation.1398 (param_0.4048: f32[128]) -> f32[128] { + %param_0.4048 = f32[128]{0:T(128)S(1)} parameter(0) + %constant.1744.clone.11 = f32[]{:T(128)} constant(0.001953125) + %broadcast.4436 = f32[128]{0:T(128)} broadcast(%constant.1744.clone.11), dimensions={}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call" stack_frame_id=0} + %div.2476 = f32[128]{0:T(128)} multiply(%param_0.4048, %broadcast.4436), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/div" stack_frame_id=0} + %constant.1745.clone.14 = f32[]{:T(128)} constant(1e-06) + %broadcast.4435 = f32[128]{0:T(128)} broadcast(%constant.1745.clone.14), dimensions={}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call" stack_frame_id=0} + %add.3659 = f32[128]{0:T(128)} add(%div.2476, %broadcast.4435), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/add" stack_frame_id=0} + ROOT %rsqrt.345 = f32[128]{0:T(128)S(1)} rsqrt(%add.3659), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/rsqrt" stack_frame_id=0} +} + +%convert_element_type.2152.reduce_sub_computation (lhs.21: bf16[], rhs.21: bf16[]) -> bf16[] { + %lhs.21 = bf16[] parameter(0) + %rhs.21 = bf16[] parameter(1) + ROOT %add.3053 = bf16[] add(%lhs.21, %rhs.21), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%fused_computation.1186 (param_0.4034: f32[512,1]) -> bf16[512] { + %param_0.4034 = f32[512,1]{0,1:T(1,128)} parameter(0) + %convert_element_type.2835 = bf16[512,1]{0,1:T(2,128)(2,1)} convert(%param_0.4034), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %constant.4132 = bf16[]{:T(256)} constant(-0), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + ROOT %reduce.803 = bf16[512]{0:T(512)(128)(2,1)S(1)} reduce(%convert_element_type.2835, %constant.4132), dimensions={1}, to_apply=%convert_element_type.2152.reduce_sub_computation, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} +} + +%fused_computation.949.clone.1.clone.clone (param_0.4341: bf16[1,128,512], param_1.5117: f32[128], param_2.4113: bf16[512]) -> bf16[128,512] { + %param_2.4113 = bf16[512]{0:T(512)(128)(2,1)S(1)} parameter(2) + %dot_general.838 = bf16[1,128,512]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.4113), dimensions={2}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/...k,k->...k/dot_general" stack_frame_id=0} + %convert.840 = f32[1,128,512]{2,1,0:T(8,128)} convert(%dot_general.838) + %param_0.4341 = bf16[1,128,512]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.2904 = f32[1,128,512]{2,1,0:T(8,128)} convert(%param_0.4341), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %param_1.5117 = f32[128]{0:T(128)S(1)} parameter(1) + %mul.4918 = f32[1,128,512]{2,1,0:T(8,128)} broadcast(%param_1.5117), dimensions={1}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %mul.4917 = f32[1,128,512]{2,1,0:T(8,128)} multiply(%convert_element_type.2904, %mul.4918), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %convert_element_type.2903 = bf16[1,128,512]{2,1,0:T(8,128)(2,1)} convert(%mul.4917), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %convert.841 = f32[1,128,512]{2,1,0:T(8,128)} convert(%convert_element_type.2903) + %dot_general.837 = f32[1,128,512]{2,1,0:T(8,128)} multiply(%convert.840, %convert.841), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/...k,k->...k/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","512"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %convert.842 = bf16[1,128,512]{2,1,0:T(8,128)(2,1)} convert(%dot_general.837) + ROOT %bitcast.2109 = bf16[128,512]{1,0:T(8,128)(2,1)} bitcast(%convert.842), metadata={stack_frame_id=0} +} + +%fused_computation.731.clone.clone.clone.clone.clone.clone (param_0.4340: bf16[512,1,1536]) -> bf16[512,1536] { + %param_0.4340 = bf16[512,1,1536]{2,0,1:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.2108 = bf16[512,1536]{1,0:T(8,128)(2,1)} bitcast(%param_0.4340), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} +} + +%region_9.14 (reduce_sum.40: f32[], reduce_sum.41: f32[]) -> f32[] { + %reduce_sum.40 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/reduce_sum"} + %reduce_sum.41 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/reduce_sum"} + ROOT %reduce_sum.45 = f32[]{:T(128)} add(%reduce_sum.40, %reduce_sum.41), metadata={op_name="reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%fused_computation.766 (param_0.4493: bf16[512,1,1536], param_1.5223: bf16[1,128,512], param_2.4190: f32[128], param_3.2841: bf16[512]) -> (f32[128], f32[1,128,1536]) { + %param_1.5223 = bf16[1,128,512]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %param_2.4190 = f32[128]{0:T(128)S(1)} parameter(2) + %param_3.2841 = bf16[512]{0:T(512)(128)(2,1)S(1)} parameter(3) + %fusion.717.clone.1 = bf16[128,512]{1,0:T(8,128)(2,1)} fusion(%param_1.5223, %param_2.4190, %param_3.2841), kind=kLoop, calls=%fused_computation.949.clone.1.clone.clone, metadata={stack_frame_id=0} + %param_0.4493 = bf16[512,1,1536]{2,0,1:T(8,128)(2,1)S(1)} parameter(0) + %fusion.481.clone.1 = bf16[512,1536]{1,0:T(8,128)(2,1)} fusion(%param_0.4493), kind=kLoop, calls=%fused_computation.731.clone.clone.clone.clone.clone.clone, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} + %convolution.191.clone.1 = bf16[128,1536]{1,0:T(8,128)(2,1)} convolution(%fusion.717.clone.1, %fusion.481.clone.1), dim_labels=bf_io->bf, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} + %convert_element_type.2617.clone.1 = f32[128,1536]{1,0:T(8,128)} convert(%convolution.191.clone.1), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %bitcast.1485.clone.1 = f32[1,128,1536]{2,1,0:T(8,128)S(1)} bitcast(%convert_element_type.2617.clone.1), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %square.609 = f32[1,128,1536]{2,1,0:T(8,128)} multiply(%bitcast.1485.clone.1, %bitcast.1485.clone.1), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/square" stack_frame_id=0} + %constant.1730.clone.93 = f32[]{:T(128)} constant(0) + %reduce.749 = f32[128]{0:T(128)S(1)} reduce(%square.609, %constant.1730.clone.93), dimensions={0,2}, to_apply=%region_9.14, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/reduce_sum" stack_frame_id=0} + ROOT %tuple.725 = (f32[128]{0:T(128)S(1)}, f32[1,128,1536]{2,1,0:T(8,128)S(1)}) tuple(%reduce.749, %bitcast.1485.clone.1) +} + +%fused_computation.1373 (param_0.3997: f32[128]) -> f32[128] { + %param_0.3997 = f32[128]{0:T(128)S(1)} parameter(0) + %constant.1674.clone.4 = f32[]{:T(128)} constant(0.000651041686), metadata={stack_frame_id=0} + %broadcast.4382 = f32[128]{0:T(128)} broadcast(%constant.1674.clone.4), dimensions={}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/div" stack_frame_id=0} + %div.2431 = f32[128]{0:T(128)} multiply(%param_0.3997, %broadcast.4382), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/div" stack_frame_id=0} + %constant.1745.clone.3 = f32[]{:T(128)} constant(1e-06) + %broadcast.4393 = f32[128]{0:T(128)} broadcast(%constant.1745.clone.3), dimensions={}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call" stack_frame_id=0} + %add.3619 = f32[128]{0:T(128)} add(%div.2431, %broadcast.4393), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/add" stack_frame_id=0} + ROOT %rsqrt.334 = f32[128]{0:T(128)S(1)} rsqrt(%add.3619), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/rsqrt" stack_frame_id=0} +} + +%convert_element_type.2151.reduce_sub_computation (lhs.20: bf16[], rhs.20: bf16[]) -> bf16[] { + %lhs.20 = bf16[] parameter(0) + %rhs.20 = bf16[] parameter(1) + ROOT %add.3052 = bf16[] add(%lhs.20, %rhs.20), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%fused_computation.1106 (param_0.4035: f32[1536,1]) -> bf16[1536] { + %param_0.4035 = f32[1536,1]{0,1:T(1,128)S(1)} parameter(0) + %convert_element_type.2829 = bf16[1536,1]{0,1:T(2,128)(2,1)} convert(%param_0.4035), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %constant.4133 = bf16[]{:T(256)} constant(-0), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + ROOT %reduce.794 = bf16[1536]{0:T(1024)(128)(2,1)S(1)} reduce(%convert_element_type.2829, %constant.4133), dimensions={1}, to_apply=%convert_element_type.2151.reduce_sub_computation, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} +} + +%fused_computation.845.clone (param_0.2713: f32[1,128,1536], param_1.2956: f32[128], param_2.2081: bf16[1536]) -> bf16[128,1536,1] { + %param_2.2081 = bf16[1536]{0:T(1024)(128)(2,1)S(1)} parameter(2) + %dot_general.739 = bf16[1,128,1536]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.2081), dimensions={2}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/...k,k->...k/dot_general" stack_frame_id=0} + %convert.837 = f32[1,128,1536]{2,1,0:T(8,128)} convert(%dot_general.739) + %param_0.2713 = f32[1,128,1536]{2,1,0:T(8,128)S(1)} parameter(0) + %param_1.2956 = f32[128]{0:T(128)S(1)} parameter(1) + %mul.4053 = f32[1,128,1536]{2,1,0:T(8,128)} broadcast(%param_1.2956), dimensions={1}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %mul.4049 = f32[1,128,1536]{2,1,0:T(8,128)} multiply(%param_0.2713, %mul.4053), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %convert_element_type.2648 = bf16[1,128,1536]{2,1,0:T(8,128)(2,1)} convert(%mul.4049), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %convert.838 = f32[1,128,1536]{2,1,0:T(8,128)} convert(%convert_element_type.2648) + %dot_general.718 = f32[1,128,1536]{2,1,0:T(8,128)} multiply(%convert.837, %convert.838), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/...k,k->...k/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","1536"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %convert.839 = bf16[1,128,1536]{2,1,0:T(8,128)(2,1)} convert(%dot_general.718) + ROOT %bitcast.1630 = bf16[128,1536,1]{1,0,2:T(8,128)(2,1)} bitcast(%convert.839) +} + +%bitcast_fusion.12 (bitcast_input.12: bf16[1536,128,192]) -> bf16[1536,128,192] { + %bitcast_input.12 = bf16[1536,128,192]{0,2,1:T(8,128)(2,1)} parameter(0) + ROOT %bitcast.2170 = bf16[1536,128,192]{0,2,1:T(8,128)(2,1)} bitcast(%bitcast_input.12) +} + +%fused_computation.700 (param_0.2453: bf16[1536,128,192], param_1.2955: f32[1,128,1536], param_2.2080: f32[128], param_3.1199: bf16[1536]) -> bf16[1,128,128,192] { + %param_1.2955 = f32[1,128,1536]{2,1,0:T(8,128)S(1)} parameter(1) + %param_2.2080 = f32[128]{0:T(128)S(1)} parameter(2) + %param_3.1199 = bf16[1536]{0:T(1024)(128)(2,1)S(1)} parameter(3) + %fusion.614 = bf16[128,1536,1]{1,0,2:T(8,128)(2,1)} fusion(%param_1.2955, %param_2.2080, %param_3.1199), kind=kLoop, calls=%fused_computation.845.clone + %param_0.2453 = bf16[1536,128,192]{0,2,1:T(8,128)(2,1)} parameter(0) + %fusion.903 = bf16[1536,128,192]{0,2,1:T(8,128)(2,1)} fusion(%param_0.2453), kind=kLoop, calls=%bitcast_fusion.12 + %convolution.183 = bf16[128,128,192]{2,0,1:T(8,128)(2,1)} convolution(%fusion.614, %fusion.903), window={size=128 pad=127_127 rhs_reversal=1}, dim_labels=bf0_i0o->b0f, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} + ROOT %bitcast.1436 = bf16[1,128,128,192]{3,1,2,0:T(8,128)(2,1)S(1)} bitcast(%convolution.183), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} +} + +%fused_computation.1427 (param_0.4318: bf16[1,128,128,32,2]) -> (bf16[1,128,128,32,1], bf16[1,128,128,32,1]) { + %param_0.4318 = bf16[1,128,128,32,2]{2,1,4,3,0:T(8,128)(2,1)S(1)} parameter(0) + %slice.1337 = bf16[1,128,128,32,1]{2,1,4,3,0:T(8,128)(2,1)S(1)} slice(%param_0.4318), slice={[0:1], [0:128], [0:128], [0:32], [0:1]}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/slice" stack_frame_id=0} + %slice.1338 = bf16[1,128,128,32,1]{2,1,4,3,0:T(8,128)(2,1)S(1)} slice(%param_0.4318), slice={[0:1], [0:128], [0:128], [0:32], [1:2]}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/slice" stack_frame_id=0} + ROOT %tuple.679 = (bf16[1,128,128,32,1]{2,1,4,3,0:T(8,128)(2,1)S(1)}, bf16[1,128,128,32,1]{2,1,4,3,0:T(8,128)(2,1)S(1)}) tuple(%slice.1337, %slice.1338) +} + +%fused_computation.685 (param_0.2073: bf16[1,128,128,32,1], param_1.2332: bf16[1,128,128,32,1], param_2.4189: f32[128,32], param_3.2840: f32[128,32]) -> (bf16[1,128,128,32], bf16[1,128,128,32]) { + %param_1.2332 = bf16[1,128,128,32,1]{3,1,2,4,0:T(8,128)(2,1)S(1)} parameter(1) + %convert_element_type.2591 = f32[1,128,128,32,1]{3,1,2,4,0:T(8,128)} convert(%param_1.2332), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %bitcast.1411 = f32[1,128,128,32]{3,1,2,0:T(8,128)} bitcast(%convert_element_type.2591), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %param_0.2073 = bf16[1,128,128,32,1]{3,1,2,4,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.2590 = f32[1,128,128,32,1]{3,1,2,4,0:T(8,128)} convert(%param_0.2073), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %bitcast.1410 = f32[1,128,128,32]{3,1,2,0:T(8,128)} bitcast(%convert_element_type.2590), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %constant.1730.clone.92 = f32[]{:T(128)} constant(0) + %convert_element_type.2612 = f32[1,128,128,32]{3,1,2,0:T(8,128)} broadcast(%constant.1730.clone.92), dimensions={}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %mul.3893 = f32[1,128,128,32]{3,1,2,0:T(8,128)} multiply(%bitcast.1410, %convert_element_type.2612), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %add.3320 = f32[1,128,128,32]{3,1,2,0:T(8,128)} add(%bitcast.1411, %mul.3893), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/add" stack_frame_id=0} + %param_3.2840 = f32[128,32]{1,0:T(8,128)S(1)} parameter(3) + %broadcast_in_dim.1473 = f32[1,128,128,32]{3,1,2,0:T(8,128)} broadcast(%param_3.2840), dimensions={1,3}, metadata={op_name="jit(train_step)/broadcast_in_dim" stack_frame_id=0} + %mul.3889 = f32[1,128,128,32]{3,1,2,0:T(8,128)} multiply(%add.3320, %broadcast_in_dim.1473), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %param_2.4189 = f32[128,32]{1,0:T(8,128)S(1)} parameter(2) + %broadcast_in_dim.1471 = f32[1,128,128,32]{3,1,2,0:T(8,128)} broadcast(%param_2.4189), dimensions={1,3}, metadata={op_name="jit(train_step)/broadcast_in_dim" stack_frame_id=0} + %mul.3880 = f32[1,128,128,32]{3,1,2,0:T(8,128)} multiply(%bitcast.1410, %broadcast_in_dim.1471), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %mul.3879 = f32[1,128,128,32]{3,1,2,0:T(8,128)} subtract(%mul.3889, %mul.3880), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %convert.529 = bf16[1,128,128,32]{3,1,2,0:T(8,128)(2,1)S(1)} convert(%mul.3879), metadata={op_name="convert.350" stack_frame_id=0} + %mul.3888.clone.1 = f32[1,128,128,32]{3,1,2,0:T(8,128)} multiply(%add.3320, %broadcast_in_dim.1471), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %mul.3887.clone.1 = f32[1,128,128,32]{3,1,2,0:T(8,128)} multiply(%bitcast.1410, %broadcast_in_dim.1473), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %mul.3886.clone.1 = f32[1,128,128,32]{3,1,2,0:T(8,128)} add(%mul.3888.clone.1, %mul.3887.clone.1), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %convert.530.clone.1 = bf16[1,128,128,32]{3,1,2,0:T(8,128)(2,1)S(1)} convert(%mul.3886.clone.1), metadata={op_name="convert.351" stack_frame_id=0} + ROOT %tuple.668 = (bf16[1,128,128,32]{3,1,2,0:T(8,128)(2,1)S(1)}, bf16[1,128,128,32]{3,1,2,0:T(8,128)(2,1)S(1)}) tuple(%convert.529, %convert.530.clone.1) +} + +%fused_computation.654 (param_0.2068: bf16[1,128,128,32], param_1.4738: bf16[1,128,128,32], param_2.3577: bf16[1,128,128,128]) -> bf16[128,128,192] { + %param_2.3577 = bf16[1,128,128,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(2) + %constant.1746.clone.13 = bf16[]{:T(256)} constant(-inf) + %pad.283 = bf16[1,128,128,192]{3,1,2,0:T(8,128)(2,1)} pad(%param_2.3577, %constant.1746.clone.13), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/concatenate" stack_frame_id=0} + %convert.676 = f32[1,128,128,192]{3,1,2,0:T(8,128)} convert(%pad.283) + %param_1.4738 = bf16[1,128,128,32]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) + %pad.282 = bf16[1,128,128,192]{3,1,2,0:T(8,128)(2,1)} pad(%param_1.4738, %constant.1746.clone.13), padding=0_0x0_0x0_0x128_32, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/concatenate" stack_frame_id=0} + %convert.677 = f32[1,128,128,192]{3,1,2,0:T(8,128)} convert(%pad.282) + %maximum.48 = f32[1,128,128,192]{3,1,2,0:T(8,128)} maximum(%convert.676, %convert.677), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/concatenate" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","128","192"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %param_0.2068 = bf16[1,128,128,32]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %pad.281 = bf16[1,128,128,192]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.2068, %constant.1746.clone.13), padding=0_0x0_0x0_0x160_0, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/concatenate" stack_frame_id=0} + %convert.678 = f32[1,128,128,192]{3,1,2,0:T(8,128)} convert(%pad.281) + %maximum.47 = f32[1,128,128,192]{3,1,2,0:T(8,128)} maximum(%maximum.48, %convert.678), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/concatenate" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","128","192"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %constant.1684.clone.4 = bf16[]{:T(256)} constant(0.1348), metadata={stack_frame_id=0} + %mul.3950 = bf16[1,128,128,192]{3,1,2,0:T(8,128)(2,1)} broadcast(%constant.1684.clone.4), dimensions={}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %convert.679 = f32[1,128,128,192]{3,1,2,0:T(8,128)} convert(%mul.3950) + %mul.3830 = f32[1,128,128,192]{3,1,2,0:T(8,128)} multiply(%maximum.47, %convert.679), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","128","192"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %convert.680 = bf16[1,128,128,192]{3,1,2,0:T(8,128)(2,1)} convert(%mul.3830) + ROOT %bitcast.1352 = bf16[128,128,192]{2,1,0:T(8,128)(2,1)} bitcast(%convert.680), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/vmap(static_sparsity_splash)/jit(_splash_attention)/splash_mha_fwd_segmented_residuals/splash_mha_fwd_segmented_residuals/squeeze" stack_frame_id=0} +} + +%fused_computation.949.clone (param_0.3007: bf16[1,128,512], param_1.3340: f32[128], param_2.2345: bf16[512]) -> bf16[128,512] { + %param_2.2345 = bf16[512]{0:T(512)(128)(2,1)S(1)} parameter(2) + %dot_general.799 = bf16[1,128,512]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.2345), dimensions={2}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/...k,k->...k/dot_general" stack_frame_id=0} + %convert.843 = f32[1,128,512]{2,1,0:T(8,128)} convert(%dot_general.799) + %param_0.3007 = bf16[1,128,512]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.2787 = f32[1,128,512]{2,1,0:T(8,128)} convert(%param_0.3007), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %param_1.3340 = f32[128]{0:T(128)S(1)} parameter(1) + %mul.4257 = f32[1,128,512]{2,1,0:T(8,128)} broadcast(%param_1.3340), dimensions={1}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %mul.4256 = f32[1,128,512]{2,1,0:T(8,128)} multiply(%convert_element_type.2787, %mul.4257), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %convert_element_type.2786 = bf16[1,128,512]{2,1,0:T(8,128)(2,1)} convert(%mul.4256), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %convert.844 = f32[1,128,512]{2,1,0:T(8,128)} convert(%convert_element_type.2786) + %dot_general.780 = f32[1,128,512]{2,1,0:T(8,128)} multiply(%convert.843, %convert.844), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/...k,k->...k/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","512"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %convert.845 = bf16[1,128,512]{2,1,0:T(8,128)(2,1)} convert(%dot_general.780) + ROOT %bitcast.1753 = bf16[128,512]{1,0:T(8,128)(2,1)} bitcast(%convert.845), metadata={stack_frame_id=0} +} + +%fused_computation.780.clone.clone (param_0.2255: bf16[512,1,576]) -> bf16[512,576] { + %param_0.2255 = bf16[512,1,576]{0,2,1:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.1526 = bf16[512,576]{0,1:T(8,128)(2,1)} bitcast(%param_0.2255), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} +} + +%fused_computation.782 (param_0.2896: bf16[512,1,576], param_1.3339: bf16[1,128,512], param_2.2344: f32[128], param_3.1373: bf16[512]) -> bf16[1,128,576] { + %param_1.3339 = bf16[1,128,512]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %param_2.2344 = f32[128]{0:T(128)S(1)} parameter(2) + %param_3.1373 = bf16[512]{0:T(512)(128)(2,1)S(1)} parameter(3) + %fusion.716 = bf16[128,512]{1,0:T(8,128)(2,1)} fusion(%param_1.3339, %param_2.2344, %param_3.1373), kind=kLoop, calls=%fused_computation.949.clone, metadata={stack_frame_id=0} + %param_0.2896 = bf16[512,1,576]{0,2,1:T(8,128)(2,1)S(1)} parameter(0) + %fusion.512 = bf16[512,576]{0,1:T(8,128)(2,1)} fusion(%param_0.2896), kind=kLoop, calls=%fused_computation.780.clone.clone, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} + %convolution.207 = bf16[128,576]{0,1:T(8,128)(2,1)} convolution(%fusion.716, %fusion.512), dim_labels=bf_io->bf, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} + ROOT %bitcast.1524 = bf16[1,128,576]{1,2,0:T(8,128)(2,1)S(1)} bitcast(%convolution.207), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} +} + +%fused_computation.1430 (param_0.4321: bf16[1,128,1,32,2]) -> (bf16[1,128,1,32,1], bf16[1,128,1,32,1]) { + %param_0.4321 = bf16[1,128,1,32,2]{1,4,3,2,0:T(2,128)(2,1)S(1)} parameter(0) + %slice.1341 = bf16[1,128,1,32,1]{1,4,3,2,0:T(2,128)(2,1)S(1)} slice(%param_0.4321), slice={[0:1], [0:128], [0:1], [0:32], [0:1]}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/slice" stack_frame_id=0} + %slice.1342 = bf16[1,128,1,32,1]{1,4,3,2,0:T(2,128)(2,1)S(1)} slice(%param_0.4321), slice={[0:1], [0:128], [0:1], [0:32], [1:2]}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/slice" stack_frame_id=0} + ROOT %tuple.697 = (bf16[1,128,1,32,1]{1,4,3,2,0:T(2,128)(2,1)S(1)}, bf16[1,128,1,32,1]{1,4,3,2,0:T(2,128)(2,1)S(1)}) tuple(%slice.1341, %slice.1342) +} + +%fused_computation.1037 (param_0.2976: f32[128,32], param_1.3320: f32[128,32], param_2.2329: bf16[1,128,1,32,1], param_3.1433: bf16[1,128,1,32,1]) -> (bf16[1,128,32], bf16[1,128,32]) { + %param_3.1433 = bf16[1,128,1,32,1]{3,1,4,2,0:T(8,128)(2,1)S(1)} parameter(3) + %convert_element_type.2812 = f32[1,128,1,32,1]{3,1,4,2,0:T(8,128)} convert(%param_3.1433), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %bitcast.1829 = f32[1,128,32]{2,1,0:T(8,128)} bitcast(%convert_element_type.2812), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %param_2.2329 = bf16[1,128,1,32,1]{3,1,4,2,0:T(8,128)(2,1)S(1)} parameter(2) + %convert_element_type.2810 = f32[1,128,1,32,1]{3,1,4,2,0:T(8,128)} convert(%param_2.2329), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %bitcast.1827 = f32[1,128,32]{2,1,0:T(8,128)} bitcast(%convert_element_type.2810), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %constant.1730.clone.91 = f32[]{:T(128)} constant(0) + %convert_element_type.2820 = f32[1,128,32]{2,1,0:T(8,128)} broadcast(%constant.1730.clone.91), dimensions={}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %mul.4349 = f32[1,128,32]{2,1,0:T(8,128)} multiply(%bitcast.1827, %convert_element_type.2820), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %add.3432 = f32[1,128,32]{2,1,0:T(8,128)} add(%bitcast.1829, %mul.4349), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/add" stack_frame_id=0} + %param_1.3320 = f32[128,32]{1,0:T(8,128)S(1)} parameter(1) + %bitcast.1816 = f32[1,128,32]{2,1,0:T(8,128)} bitcast(%param_1.3320), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %mul.4340 = f32[1,128,32]{2,1,0:T(8,128)} multiply(%add.3432, %bitcast.1816), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %param_0.2976 = f32[128,32]{1,0:T(8,128)S(1)} parameter(0) + %bitcast.1808 = f32[1,128,32]{2,1,0:T(8,128)} bitcast(%param_0.2976), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %mul.4329 = f32[1,128,32]{2,1,0:T(8,128)} multiply(%bitcast.1827, %bitcast.1808), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %mul.4328 = f32[1,128,32]{2,1,0:T(8,128)} subtract(%mul.4340, %mul.4329), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %convert.537 = bf16[1,128,32]{2,1,0:T(8,128)(2,1)S(1)} convert(%mul.4328), metadata={op_name="convert.348" stack_frame_id=0} + %mul.4336.clone.1 = f32[1,128,32]{2,1,0:T(8,128)} multiply(%add.3432, %bitcast.1808), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %mul.4342.clone.1 = f32[1,128,32]{2,1,0:T(8,128)} multiply(%bitcast.1827, %bitcast.1816), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %mul.4335.clone.1 = f32[1,128,32]{2,1,0:T(8,128)} add(%mul.4336.clone.1, %mul.4342.clone.1), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %convert.538.clone.1 = bf16[1,128,32]{2,1,0:T(8,128)(2,1)S(1)} convert(%mul.4335.clone.1), metadata={op_name="convert.349" stack_frame_id=0} + ROOT %tuple.689 = (bf16[1,128,32]{2,1,0:T(8,128)(2,1)S(1)}, bf16[1,128,32]{2,1,0:T(8,128)(2,1)S(1)}) tuple(%convert.537, %convert.538.clone.1) +} + +%fused_computation.1064 (param_0.3059: bf16[1,128,32], param_1.4734: bf16[1,128,32]) -> bf16[128,64] { + %param_1.4734 = bf16[1,128,32]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %constant.1746.clone.6 = bf16[]{:T(256)} constant(-inf) + %pad.314 = bf16[1,128,64]{2,1,0:T(8,128)(2,1)} pad(%param_1.4734, %constant.1746.clone.6), padding=0_0x0_0x0_32, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/concatenate" stack_frame_id=0} + %convert.795 = f32[1,128,64]{2,1,0:T(8,128)} convert(%pad.314) + %param_0.3059 = bf16[1,128,32]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %pad.313 = bf16[1,128,64]{2,1,0:T(8,128)(2,1)} pad(%param_0.3059, %constant.1746.clone.6), padding=0_0x0_0x32_0, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/concatenate" stack_frame_id=0} + %convert.796 = f32[1,128,64]{2,1,0:T(8,128)} convert(%pad.313) + %maximum.64 = f32[1,128,64]{2,1,0:T(8,128)} maximum(%convert.795, %convert.796), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/concatenate" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","64"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %convert.797 = bf16[1,128,64]{2,1,0:T(8,128)(2,1)} convert(%maximum.64) + ROOT %bitcast.1839 = bf16[128,64]{1,0:T(8,128)(2,1)S(1)} bitcast(%convert.797), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/concatenate" stack_frame_id=0} +} + +%region_10.15 (reduce_sum.46: f32[], reduce_sum.47: f32[]) -> f32[] { + %reduce_sum.46 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/reduce_sum"} + %reduce_sum.47 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/reduce_sum"} + ROOT %reduce_sum.48 = f32[]{:T(128)} add(%reduce_sum.46, %reduce_sum.47), metadata={op_name="reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%fused_computation.888 (param_0.4492: bf16[1,128,576]) -> f32[128] { + %param_0.4492 = bf16[1,128,576]{1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %slice.1213 = bf16[1,128,512]{1,2,0:T(8,128)(2,1)} slice(%param_0.4492), slice={[0:1], [0:128], [0:512]}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/split" stack_frame_id=0} + %convert_element_type.2686 = f32[1,128,512]{1,2,0:T(8,128)} convert(%slice.1213), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %square.622 = f32[1,128,512]{1,2,0:T(8,128)} multiply(%convert_element_type.2686, %convert_element_type.2686), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/square" stack_frame_id=0} + %constant.1730.clone.90 = f32[]{:T(128)} constant(0) + ROOT %reduce.761 = f32[128]{0:T(128)S(1)} reduce(%square.622, %constant.1730.clone.90), dimensions={0,2}, to_apply=%region_10.15, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/reduce_sum" stack_frame_id=0} +} + +%fused_computation.1372 (param_0.3996: f32[128]) -> f32[128] { + %param_0.3996 = f32[128]{0:T(128)S(1)} parameter(0) + %constant.1744.clone.3 = f32[]{:T(128)} constant(0.001953125) + %broadcast.4402 = f32[128]{0:T(128)} broadcast(%constant.1744.clone.3), dimensions={}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call" stack_frame_id=0} + %div.2430 = f32[128]{0:T(128)} multiply(%param_0.3996, %broadcast.4402), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/div" stack_frame_id=0} + %constant.1745.clone.4 = f32[]{:T(128)} constant(1e-06) + %broadcast.4392 = f32[128]{0:T(128)} broadcast(%constant.1745.clone.4), dimensions={}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call" stack_frame_id=0} + %add.3618 = f32[128]{0:T(128)} add(%div.2430, %broadcast.4392), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/add" stack_frame_id=0} + ROOT %rsqrt.333 = f32[128]{0:T(128)S(1)} rsqrt(%add.3618), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/rsqrt" stack_frame_id=0} +} + +%convert_element_type.2155.reduce_sub_computation (lhs.22: bf16[], rhs.22: bf16[]) -> bf16[] { + %lhs.22 = bf16[] parameter(0) + %rhs.22 = bf16[] parameter(1) + ROOT %add.3054 = bf16[] add(%lhs.22, %rhs.22), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%fused_computation.1185 (param_0.4033: f32[512,1]) -> bf16[512] { + %param_0.4033 = f32[512,1]{0,1:T(1,128)} parameter(0) + %convert_element_type.2834 = bf16[512,1]{0,1:T(2,128)(2,1)} convert(%param_0.4033), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %constant.4131 = bf16[]{:T(256)} constant(-0), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + ROOT %reduce.802 = bf16[512]{0:T(512)(128)(2,1)S(1)} reduce(%convert_element_type.2834, %constant.4131), dimensions={1}, to_apply=%convert_element_type.2155.reduce_sub_computation, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} +} + +%fused_computation.953.clone (param_0.3004: f32[128], param_1.3338: bf16[1,128,576], param_2.2343: bf16[512]) -> bf16[128,512,1] { + %param_2.2343 = bf16[512]{0:T(512)(128)(2,1)S(1)} parameter(2) + %dot_general.795 = bf16[1,128,512]{1,2,0:T(8,128)(2,1)} broadcast(%param_2.2343), dimensions={2}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/...k,k->...k/dot_general" stack_frame_id=0} + %convert.828 = f32[1,128,512]{1,2,0:T(8,128)} convert(%dot_general.795) + %param_1.3338 = bf16[1,128,576]{1,2,0:T(8,128)(2,1)S(1)} parameter(1) + %slice.1212 = bf16[1,128,512]{1,2,0:T(8,128)(2,1)} slice(%param_1.3338), slice={[0:1], [0:128], [0:512]}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/split" stack_frame_id=0} + %convert_element_type.2749 = f32[1,128,512]{1,2,0:T(8,128)} convert(%slice.1212), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %param_0.3004 = f32[128]{0:T(128)S(1)} parameter(0) + %mul.4202 = f32[1,128,512]{1,2,0:T(8,128)} broadcast(%param_0.3004), dimensions={1}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %mul.4179 = f32[1,128,512]{1,2,0:T(8,128)} multiply(%convert_element_type.2749, %mul.4202), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %convert_element_type.2748 = bf16[1,128,512]{1,2,0:T(8,128)(2,1)} convert(%mul.4179), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %convert.829 = f32[1,128,512]{1,2,0:T(8,128)} convert(%convert_element_type.2748) + %dot_general.753 = f32[1,128,512]{1,2,0:T(8,128)} multiply(%convert.828, %convert.829), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/...k,k->...k/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","512"],"tuple_shapes":[],"layout":{"minor_to_major":["1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %convert.830 = bf16[1,128,512]{1,2,0:T(8,128)(2,1)} convert(%dot_general.753) + ROOT %bitcast.1708 = bf16[128,512,1]{0,1,2:T(8,128)(2,1)} bitcast(%convert.830) +} + +%fused_computation.594.clone.clone (param_0.2031: bf16[512,1,128,256]) -> bf16[512,128,256] { + %param_0.2031 = bf16[512,1,128,256]{3,0,2,1:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.1431 = bf16[512,128,256]{2,0,1:T(8,128)(2,1)} bitcast(%param_0.2031), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} +} + +%fused_computation.697 (param_0.2750: bf16[512,1,128,256], param_1.3337: f32[128], param_2.2342: bf16[1,128,576], param_3.1372: bf16[512]) -> bf16[1,128,128,256] { + %param_1.3337 = f32[128]{0:T(128)S(1)} parameter(1) + %param_2.2342 = bf16[1,128,576]{1,2,0:T(8,128)(2,1)S(1)} parameter(2) + %param_3.1372 = bf16[512]{0:T(512)(128)(2,1)S(1)} parameter(3) + %fusion.672 = bf16[128,512,1]{0,1,2:T(8,128)(2,1)} fusion(%param_1.3337, %param_2.2342, %param_3.1372), kind=kLoop, calls=%fused_computation.953.clone + %param_0.2750 = bf16[512,1,128,256]{3,0,2,1:T(8,128)(2,1)S(1)} parameter(0) + %fusion.455 = bf16[512,128,256]{2,0,1:T(8,128)(2,1)} fusion(%param_0.2750), kind=kLoop, calls=%fused_computation.594.clone.clone, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} + %convolution.179 = bf16[128,128,256]{2,0,1:T(8,128)(2,1)} convolution(%fusion.672, %fusion.455), window={size=128 pad=127_127 rhs_reversal=1}, dim_labels=bf0_i0o->b0f, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} + ROOT %bitcast.1429 = bf16[1,128,128,256]{3,1,2,0:T(8,128)(2,1)S(1)} bitcast(%convolution.179), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} +} + +%fused_computation.706 (param_0.2059: bf16[1,128,128,256]) -> (bf16[128,128,128], bf16[1,128,128,128]) { + %param_0.2059 = bf16[1,128,128,256]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %slice.1208 = bf16[1,128,128,128]{3,1,2,0:T(8,128)(2,1)} slice(%param_0.2059), slice={[0:1], [0:128], [0:128], [128:256]}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/split" stack_frame_id=0} + %bitcast.1440 = bf16[128,128,128]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%slice.1208), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/vmap(static_sparsity_splash)/jit(_splash_attention)/splash_mha_fwd_segmented_residuals/splash_mha_fwd_segmented_residuals/squeeze" stack_frame_id=0} + %slice.1334 = bf16[1,128,128,128]{3,1,2,0:T(8,128)(2,1)S(1)} slice(%param_0.2059), slice={[0:1], [0:128], [0:128], [0:128]}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/split" stack_frame_id=0} + ROOT %tuple.673 = (bf16[128,128,128]{2,1,0:T(8,128)(2,1)S(1)}, bf16[1,128,128,128]{3,1,2,0:T(8,128)(2,1)S(1)}) tuple(%bitcast.1440, %slice.1334) +} + +%fused_computation.655 (param_0.1904: bf16[1,128,128,64], param_1.4736: bf16[1,128,128,128]) -> bf16[128,128,192] { + %param_1.4736 = bf16[1,128,128,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) + %constant.1746.clone.8 = bf16[]{:T(256)} constant(-inf) + %pad.285 = bf16[1,128,128,192]{3,1,2,0:T(8,128)(2,1)} pad(%param_1.4736, %constant.1746.clone.8), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/concatenate" stack_frame_id=0} + %convert.681 = f32[1,128,128,192]{3,1,2,0:T(8,128)} convert(%pad.285) + %param_0.1904 = bf16[1,128,128,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %pad.284 = bf16[1,128,128,192]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1904, %constant.1746.clone.8), padding=0_0x0_0x0_0x128_0, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/concatenate" stack_frame_id=0} + %convert.682 = f32[1,128,128,192]{3,1,2,0:T(8,128)} convert(%pad.284) + %maximum.49 = f32[1,128,128,192]{3,1,2,0:T(8,128)} maximum(%convert.681, %convert.682), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/concatenate" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","128","192"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %convert.683 = bf16[1,128,128,192]{3,1,2,0:T(8,128)(2,1)} convert(%maximum.49) + ROOT %bitcast.1353 = bf16[128,128,192]{2,1,0:T(8,128)(2,1)} bitcast(%convert.683), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/vmap(static_sparsity_splash)/jit(_splash_attention)/splash_mha_fwd_segmented_residuals/splash_mha_fwd_segmented_residuals/squeeze" stack_frame_id=0} +} + +%fused_computation.707.clone.clone.clone.clone.clone.clone (param_0.4345: bf16[128,128,128]) -> bf16[128,128,128] { + %param_0.4345 = bf16[128,128,128]{2,1,0:T(8,128)(2,1)} parameter(0) + ROOT %bitcast.2113 = bf16[128,128,128]{2,0,1:T(8,128)(2,1)} bitcast(%param_0.4345) +} + +%fused_computation.642.clone.clone.clone.clone.clone.clone (param_0.4344: bf16[128,1,128,512]) -> bf16[128,128,512] { + %param_0.4344 = bf16[128,1,128,512]{3,2,0,1:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.2112 = bf16[128,128,512]{2,1,0:T(8,128)(2,1)} bitcast(%param_0.4344), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} +} + +%region_11.17 (reduce_sum.52: f32[], reduce_sum.53: f32[]) -> f32[] { + %reduce_sum.52 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/reduce_sum"} + %reduce_sum.53 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/reduce_sum"} + ROOT %reduce_sum.54 = f32[]{:T(128)} add(%reduce_sum.52, %reduce_sum.53), metadata={op_name="reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%fused_computation.890 (param_0.4491: bf16[1,128,512], param_1.5222: bf16[128,1,128,512], param_2.4188: bf16[128,128,128]) -> (f32[128], bf16[1,128,512]) { + %param_0.4491 = bf16[1,128,512]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert.724 = f32[1,128,512]{2,1,0:T(8,128)} convert(%param_0.4491) + %param_2.4188 = bf16[128,128,128]{2,1,0:T(8,128)(2,1)} parameter(2) + %fusion.605.clone.1 = bf16[128,128,128]{2,0,1:T(8,128)(2,1)} fusion(%param_2.4188), kind=kLoop, calls=%fused_computation.707.clone.clone.clone.clone.clone.clone + %param_1.5222 = bf16[128,1,128,512]{3,2,0,1:T(8,128)(2,1)S(1)} parameter(1) + %fusion.604.clone.1 = bf16[128,128,512]{2,1,0:T(8,128)(2,1)} fusion(%param_1.5222), kind=kLoop, calls=%fused_computation.642.clone.clone.clone.clone.clone.clone, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} + %convolution.263.clone.1 = bf16[128,512,1]{1,0,2:T(8,128)(2,1)} convolution(%fusion.605.clone.1, %fusion.604.clone.1), window={size=128}, dim_labels=b0f_0io->bf0, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} + %bitcast.1626.clone.1 = bf16[1,128,512]{2,1,0:T(8,128)(2,1)} bitcast(%convolution.263.clone.1), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} + %convert.725 = f32[1,128,512]{2,1,0:T(8,128)} convert(%bitcast.1626.clone.1) + %add.3371.clone.1 = f32[1,128,512]{2,1,0:T(8,128)} add(%convert.724, %convert.725), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/add" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","512"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %convert_element_type.2691 = f32[1,128,512]{2,1,0:T(8,128)} convert(%add.3371.clone.1), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %square.624 = f32[1,128,512]{2,1,0:T(8,128)} multiply(%convert_element_type.2691, %convert_element_type.2691), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/square" stack_frame_id=0} + %constant.1730.clone.89 = f32[]{:T(128)} constant(0) + %reduce.762 = f32[128]{0:T(128)S(1)} reduce(%square.624, %constant.1730.clone.89), dimensions={0,2}, to_apply=%region_11.17, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/reduce_sum" stack_frame_id=0} + %convert.726 = bf16[1,128,512]{2,1,0:T(8,128)(2,1)S(1)} convert(%add.3371.clone.1) + ROOT %tuple.726 = (f32[128]{0:T(128)S(1)}, bf16[1,128,512]{2,1,0:T(8,128)(2,1)S(1)}) tuple(%reduce.762, %convert.726) +} + +%convert_element_type.2160.reduce_sub_computation (lhs.23: bf16[], rhs.23: bf16[]) -> bf16[] { + %lhs.23 = bf16[] parameter(0) + %rhs.23 = bf16[] parameter(1) + ROOT %add.3055 = bf16[] add(%lhs.23, %rhs.23), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%fused_computation.1302 (param_0.4032: f32[256,1]) -> bf16[256] { + %param_0.4032 = f32[256,1]{0,1:T(1,128)} parameter(0) + %convert_element_type.2848 = bf16[256,1]{0,1:T(2,128)(2,1)} convert(%param_0.4032), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %constant.4130 = bf16[]{:T(256)} constant(-0), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + ROOT %reduce.825 = bf16[256]{0:T(256)(128)(2,1)S(1)} reduce(%convert_element_type.2848, %constant.4130), dimensions={1}, to_apply=%convert_element_type.2160.reduce_sub_computation, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} +} + +%fused_computation.1371 (param_0.3995: f32[128]) -> f32[128] { + %param_0.3995 = f32[128]{0:T(128)S(1)} parameter(0) + %constant.1744.clone.4 = f32[]{:T(128)} constant(0.001953125) + %broadcast.4401 = f32[128]{0:T(128)} broadcast(%constant.1744.clone.4), dimensions={}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call" stack_frame_id=0} + %div.2429 = f32[128]{0:T(128)} multiply(%param_0.3995, %broadcast.4401), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/div" stack_frame_id=0} + %constant.1745.clone.5 = f32[]{:T(128)} constant(1e-06) + %broadcast.4391 = f32[128]{0:T(128)} broadcast(%constant.1745.clone.5), dimensions={}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call" stack_frame_id=0} + %add.3617 = f32[128]{0:T(128)} add(%div.2429, %broadcast.4391), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/add" stack_frame_id=0} + ROOT %rsqrt.332 = f32[128]{0:T(128)S(1)} rsqrt(%add.3617), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/rsqrt" stack_frame_id=0} +} + +%convert_element_type.2150.reduce_sub_computation (lhs.19: bf16[], rhs.19: bf16[]) -> bf16[] { + %lhs.19 = bf16[] parameter(0) + %rhs.19 = bf16[] parameter(1) + ROOT %add.3051 = bf16[] add(%lhs.19, %rhs.19), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%fused_computation.1187 (param_0.4036: f32[512,1]) -> bf16[512] { + %param_0.4036 = f32[512,1]{0,1:T(1,128)S(1)} parameter(0) + %convert_element_type.2836 = bf16[512,1]{0,1:T(2,128)(2,1)} convert(%param_0.4036), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %constant.4134 = bf16[]{:T(256)} constant(-0), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + ROOT %reduce.804 = bf16[512]{0:T(512)(128)(2,1)S(1)} reduce(%convert_element_type.2836, %constant.4134), dimensions={1}, to_apply=%convert_element_type.2150.reduce_sub_computation, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} +} + +%fused_computation.951.clone.2 (param_0.4114: bf16[1,128,512], param_1.4846: f32[128], param_2.3699: bf16[512]) -> bf16[128,512] { + %param_2.3699 = bf16[512]{0:T(512)(128)(2,1)S(1)} parameter(2) + %dot_general.828 = bf16[1,128,512]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.3699), dimensions={2}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/...k,k->...k/dot_general" stack_frame_id=0} + %convert.858 = f32[1,128,512]{2,1,0:T(8,128)} convert(%dot_general.828) + %param_0.4114 = bf16[1,128,512]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.2875 = f32[1,128,512]{2,1,0:T(8,128)} convert(%param_0.4114), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %param_1.4846 = f32[128]{0:T(128)S(1)} parameter(1) + %mul.4590 = f32[1,128,512]{2,1,0:T(8,128)} broadcast(%param_1.4846), dimensions={1}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %mul.4589 = f32[1,128,512]{2,1,0:T(8,128)} multiply(%convert_element_type.2875, %mul.4590), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %convert_element_type.2874 = bf16[1,128,512]{2,1,0:T(8,128)(2,1)} convert(%mul.4589), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %convert.859 = f32[1,128,512]{2,1,0:T(8,128)} convert(%convert_element_type.2874) + %dot_general.827 = f32[1,128,512]{2,1,0:T(8,128)} multiply(%convert.858, %convert.859), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/...k,k->...k/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","512"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %convert.860 = bf16[1,128,512]{2,1,0:T(8,128)(2,1)} convert(%dot_general.827) + ROOT %bitcast.2003 = bf16[128,512]{1,0:T(8,128)(2,1)} bitcast(%convert.860), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/reshape" stack_frame_id=0} +} + +%fused_computation.907 (param_0.2625: bf16[512,1,256]) -> bf16[512,256] { + %param_0.2625 = bf16[512,1,256]{2,0,1:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.1674 = bf16[512,256]{1,0:T(8,128)(2,1)} bitcast(%param_0.2625), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} +} + +%fused_computation.838 (param_0.4115: bf16[512,1,256], param_1.4847: bf16[1,128,512], param_2.3700: f32[128], param_3.2295: bf16[512]) -> bf16[128,256] { + %constant.1687.clone.15 = bf16[]{:T(256)} constant(1), metadata={stack_frame_id=0} + %broadcast.4228 = bf16[128,256]{1,0:T(8,128)(2,1)} broadcast(%constant.1687.clone.15), dimensions={}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call" stack_frame_id=0} + %convert.709 = f32[128,256]{1,0:T(8,128)} convert(%broadcast.4228) + %param_1.4847 = bf16[1,128,512]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %param_2.3700 = f32[128]{0:T(128)S(1)} parameter(2) + %param_3.2295 = bf16[512]{0:T(512)(128)(2,1)S(1)} parameter(3) + %fusion.824 = bf16[128,512]{1,0:T(8,128)(2,1)} fusion(%param_1.4847, %param_2.3700, %param_3.2295), kind=kLoop, calls=%fused_computation.951.clone.2, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/reshape" stack_frame_id=0} + %param_0.4115 = bf16[512,1,256]{2,0,1:T(8,128)(2,1)S(1)} parameter(0) + %fusion.651 = bf16[512,256]{1,0:T(8,128)(2,1)} fusion(%param_0.4115), kind=kLoop, calls=%fused_computation.907, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} + %convolution.277 = bf16[128,256]{1,0:T(8,128)(2,1)} convolution(%fusion.824, %fusion.651), dim_labels=bf_io->bf, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} + %convert.708 = f32[128,256]{1,0:T(8,128)} convert(%convolution.277) + %neg.218 = f32[128,256]{1,0:T(8,128)} negate(%convert.708), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/neg" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["128","256"],"tuple_shapes":[],"layout":{"minor_to_major":["1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false]}}} + %exp.511 = f32[128,256]{1,0:T(8,128)} exponential(%neg.218), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/exp" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["128","256"],"tuple_shapes":[],"layout":{"minor_to_major":["1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false]}}} + %add.3377 = f32[128,256]{1,0:T(8,128)} add(%exp.511, %convert.709), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/add" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["128","256"],"tuple_shapes":[],"layout":{"minor_to_major":["1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false]}}} + %div.2162 = f32[128,256]{1,0:T(8,128)} divide(%convert.709, %add.3377), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/div" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["128","256"],"tuple_shapes":[],"layout":{"minor_to_major":["1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false]}}} + ROOT %convert.710 = bf16[128,256]{1,0:T(8,128)(2,1)S(1)} convert(%div.2162) +} + +%fused_computation.1334 (param_0.4047: f32[128]) -> f32[128] { + %param_0.4047 = f32[128]{0:T(128)S(1)} parameter(0) + %constant.1744.clone.1 = f32[]{:T(128)} constant(0.001953125) + %broadcast.4432 = f32[128]{0:T(128)} broadcast(%constant.1744.clone.1), dimensions={}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call" stack_frame_id=0} + %div.2474 = f32[128]{0:T(128)} multiply(%param_0.4047, %broadcast.4432), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/div" stack_frame_id=0} + %constant.1745.clone.1 = f32[]{:T(128)} constant(1e-06) + %broadcast.4431 = f32[128]{0:T(128)} broadcast(%constant.1745.clone.1), dimensions={}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call" stack_frame_id=0} + %add.3657 = f32[128]{0:T(128)} add(%div.2474, %broadcast.4431), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/add" stack_frame_id=0} + %bitcast.1956 = f32[1,128]{1,0:T(1,128)} bitcast(%add.3657), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/add" stack_frame_id=0} + %rsqrt.309 = f32[1,128]{1,0:T(1,128)} rsqrt(%bitcast.1956), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/rsqrt" stack_frame_id=0} + ROOT %bitcast.1925 = f32[128]{0:T(128)S(1)} bitcast(%rsqrt.309), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/rsqrt" stack_frame_id=0} +} + +%fused_computation.1009 (param_0.3008: f32[128], param_1.3341: bf16[1,128,512], param_2.2346: bf16[512]) -> bf16[1,128,512] { + %param_2.2346 = bf16[512]{0:T(512)(128)(2,1)S(1)} parameter(2) + %dot_general.802 = bf16[1,128,512]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.2346), dimensions={2}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/...k,k->...k/dot_general" stack_frame_id=0} + %convert.775 = f32[1,128,512]{2,1,0:T(8,128)} convert(%dot_general.802) + %param_1.3341 = bf16[1,128,512]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %convert_element_type.2785 = f32[1,128,512]{2,1,0:T(8,128)} convert(%param_1.3341), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %param_0.3008 = f32[128]{0:T(128)S(1)} parameter(0) + %mul.4255 = f32[1,128,512]{2,1,0:T(8,128)} broadcast(%param_0.3008), dimensions={1}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %mul.4254 = f32[1,128,512]{2,1,0:T(8,128)} multiply(%convert_element_type.2785, %mul.4255), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %convert_element_type.2784 = bf16[1,128,512]{2,1,0:T(8,128)(2,1)} convert(%mul.4254), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %convert.776 = f32[1,128,512]{2,1,0:T(8,128)} convert(%convert_element_type.2784) + %dot_general.779 = f32[1,128,512]{2,1,0:T(8,128)} multiply(%convert.775, %convert.776), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/...k,k->...k/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","512"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + ROOT %convert.777 = bf16[1,128,512]{2,1,0:T(8,128)(2,1)S(1)} convert(%dot_general.779) +} + +%fused_computation.974.clone.clone (param_0.4333: bf16[1,128,512]) -> bf16[128,512] { + %param_0.4333 = bf16[1,128,512]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.2101 = bf16[128,512]{1,0:T(8,128)(2,1)} bitcast(%param_0.4333), metadata={stack_frame_id=0} +} + +%fused_computation.729.clone.clone.clone.clone.clone.clone (param_0.4332: bf16[512,1,1536]) -> bf16[512,1536] { + %param_0.4332 = bf16[512,1,1536]{2,0,1:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.2100 = bf16[512,1536]{1,0:T(8,128)(2,1)} bitcast(%param_0.4332), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} +} + +%region_43.52 (reduce_sum.104: f32[], reduce_sum.108: f32[]) -> f32[] { + %reduce_sum.104 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/reduce_sum"} + %reduce_sum.108 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/reduce_sum"} + ROOT %reduce_sum.109 = f32[]{:T(128)} add(%reduce_sum.104, %reduce_sum.108), metadata={op_name="checkpoint/rematted_computation/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%fused_computation.769 (param_0.4490: bf16[512,1,1536], param_1.5221: bf16[1,128,512]) -> (f32[128], f32[1,128,1536]) { + %param_1.5221 = bf16[1,128,512]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %fusion.691.clone.1 = bf16[128,512]{1,0:T(8,128)(2,1)} fusion(%param_1.5221), kind=kLoop, calls=%fused_computation.974.clone.clone, metadata={stack_frame_id=0} + %param_0.4490 = bf16[512,1,1536]{2,0,1:T(8,128)(2,1)S(1)} parameter(0) + %fusion.487.clone.1 = bf16[512,1536]{1,0:T(8,128)(2,1)} fusion(%param_0.4490), kind=kLoop, calls=%fused_computation.729.clone.clone.clone.clone.clone.clone, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} + %convolution.195.clone.1 = bf16[128,1536]{1,0:T(8,128)(2,1)} convolution(%fusion.691.clone.1, %fusion.487.clone.1), dim_labels=bf_io->bf, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/dot_general" stack_frame_id=0} + %convert_element_type.2620.clone.1 = f32[128,1536]{1,0:T(8,128)} convert(%convolution.195.clone.1), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %bitcast.1490.clone.1 = f32[1,128,1536]{2,1,0:T(8,128)S(1)} bitcast(%convert_element_type.2620.clone.1), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %square.610 = f32[1,128,1536]{2,1,0:T(8,128)} multiply(%bitcast.1490.clone.1, %bitcast.1490.clone.1), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/square" stack_frame_id=0} + %constant.1730.clone.88 = f32[]{:T(128)} constant(0) + %reduce.751 = f32[128]{0:T(128)S(1)} reduce(%square.610, %constant.1730.clone.88), dimensions={0,2}, to_apply=%region_43.52, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/reduce_sum" stack_frame_id=0} + ROOT %tuple.723 = (f32[128]{0:T(128)S(1)}, f32[1,128,1536]{2,1,0:T(8,128)S(1)}) tuple(%reduce.751, %bitcast.1490.clone.1) +} + +%fused_computation.1332 (param_0.3992: f32[128]) -> f32[128] { + %param_0.3992 = f32[128]{0:T(128)S(1)} parameter(0) + %constant.1674.clone.2 = f32[]{:T(128)} constant(0.000651041686), metadata={stack_frame_id=0} + %broadcast.4381 = f32[128]{0:T(128)} broadcast(%constant.1674.clone.2), dimensions={}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/div" stack_frame_id=0} + %div.2461 = f32[128]{0:T(128)} multiply(%param_0.3992, %broadcast.4381), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/div" stack_frame_id=0} + %constant.1745.clone.8 = f32[]{:T(128)} constant(1e-06) + %broadcast.4388 = f32[128]{0:T(128)} broadcast(%constant.1745.clone.8), dimensions={}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call" stack_frame_id=0} + %add.3640 = f32[128]{0:T(128)} add(%div.2461, %broadcast.4388), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/add" stack_frame_id=0} + %bitcast.1954 = f32[1,128]{1,0:T(1,128)} bitcast(%add.3640), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/add" stack_frame_id=0} + %rsqrt.307 = f32[1,128]{1,0:T(1,128)} rsqrt(%bitcast.1954), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/rsqrt" stack_frame_id=0} + ROOT %bitcast.1924 = f32[128]{0:T(128)S(1)} bitcast(%rsqrt.307), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/rsqrt" stack_frame_id=0} +} + +%fused_computation.768 (param_0.2222: f32[1,128,1536], param_1.2705: f32[128]) -> bf16[1,128,1536] { + %param_0.2222 = f32[1,128,1536]{2,1,0:T(8,128)S(1)} parameter(0) + %param_1.2705 = f32[128]{0:T(128)S(1)} parameter(1) + %mul.4056 = f32[1,128,1536]{2,1,0:T(8,128)} broadcast(%param_1.2705), dimensions={1}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %mul.4026 = f32[1,128,1536]{2,1,0:T(8,128)} multiply(%param_0.2222, %mul.4056), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + ROOT %convert_element_type.2631 = bf16[1,128,1536]{2,1,0:T(8,128)(2,1)S(1)} convert(%mul.4026), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} +} + +%fused_computation.853 (param_0.2710: bf16[1,128,1536], param_1.2954: bf16[1536]) -> bf16[128,1536,1] { + %param_1.2954 = bf16[1536]{0:T(1024)(128)(2,1)S(1)} parameter(1) + %dot_general.735 = bf16[1,128,1536]{2,1,0:T(8,128)(2,1)} broadcast(%param_1.2954), dimensions={2}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/...k,k->...k/dot_general" stack_frame_id=0} + %convert.834 = f32[1,128,1536]{2,1,0:T(8,128)} convert(%dot_general.735) + %param_0.2710 = bf16[1,128,1536]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert.835 = f32[1,128,1536]{2,1,0:T(8,128)} convert(%param_0.2710) + %dot_general.727 = f32[1,128,1536]{2,1,0:T(8,128)} multiply(%convert.834, %convert.835), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/...k,k->...k/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","1536"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %convert.836 = bf16[1,128,1536]{2,1,0:T(8,128)(2,1)} convert(%dot_general.727) + ROOT %bitcast.1640 = bf16[128,1536,1]{1,0,2:T(8,128)(2,1)} bitcast(%convert.836) +} + +%bitcast_fusion.11 (bitcast_input.11: bf16[1536,128,192]) -> bf16[1536,128,192] { + %bitcast_input.11 = bf16[1536,128,192]{0,2,1:T(8,128)(2,1)} parameter(0) + ROOT %bitcast.2169 = bf16[1536,128,192]{0,2,1:T(8,128)(2,1)} bitcast(%bitcast_input.11) +} + +%fused_computation.699 (param_0.2466: bf16[1536,128,192], param_1.2953: bf16[1,128,1536], param_2.2079: bf16[1536]) -> bf16[1,128,128,192] { + %param_1.2953 = bf16[1,128,1536]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %param_2.2079 = bf16[1536]{0:T(1024)(128)(2,1)S(1)} parameter(2) + %fusion.618 = bf16[128,1536,1]{1,0,2:T(8,128)(2,1)} fusion(%param_1.2953, %param_2.2079), kind=kLoop, calls=%fused_computation.853 + %param_0.2466 = bf16[1536,128,192]{0,2,1:T(8,128)(2,1)} parameter(0) + %fusion.902 = bf16[1536,128,192]{0,2,1:T(8,128)(2,1)} fusion(%param_0.2466), kind=kLoop, calls=%bitcast_fusion.11 + %convolution.182 = bf16[128,128,192]{2,0,1:T(8,128)(2,1)} convolution(%fusion.618, %fusion.902), window={size=128 pad=127_127 rhs_reversal=1}, dim_labels=bf0_i0o->b0f, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/dot_general" stack_frame_id=0} + ROOT %bitcast.1435 = bf16[1,128,128,192]{3,1,2,0:T(8,128)(2,1)S(1)} bitcast(%convolution.182), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/dot_general" stack_frame_id=0} +} + +%fused_computation.1426 (param_0.4317: bf16[1,128,128,32,2]) -> (bf16[1,128,128,32,1], bf16[1,128,128,32,1]) { + %param_0.4317 = bf16[1,128,128,32,2]{2,1,4,3,0:T(8,128)(2,1)S(1)} parameter(0) + %slice.1335 = bf16[1,128,128,32,1]{2,1,4,3,0:T(8,128)(2,1)} slice(%param_0.4317), slice={[0:1], [0:128], [0:128], [0:32], [0:1]}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/slice" stack_frame_id=0} + %slice.1336 = bf16[1,128,128,32,1]{2,1,4,3,0:T(8,128)(2,1)} slice(%param_0.4317), slice={[0:1], [0:128], [0:128], [0:32], [1:2]}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/slice" stack_frame_id=0} + ROOT %tuple.678 = (bf16[1,128,128,32,1]{2,1,4,3,0:T(8,128)(2,1)}, bf16[1,128,128,32,1]{2,1,4,3,0:T(8,128)(2,1)}) tuple(%slice.1335, %slice.1336) +} + +%fused_computation.1357 (param_0.4009: s32[1,128]) -> s32[128] { + %param_0.4009 = s32[1,128]{1,0:T(1,128)S(1)} parameter(0) + %bitcast.1952 = s32[1,128,1]{1,2,0:T(1,128)} bitcast(%param_0.4009), metadata={op_name="reshape.2486" stack_frame_id=0} + %constant.1719.clone.141 = s32[]{:T(128)} constant(0), metadata={op_name="jit(train_step)/shard_map/jit(gmm)" stack_frame_id=0} + %closed_call.244 = s32[1,128,1]{1,2,0:T(1,128)} broadcast(%constant.1719.clone.141), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call" stack_frame_id=0} + %lt.836 = pred[1,128,1]{1,2,0:T(4,128)(4,1)} compare(%bitcast.1952, %closed_call.244), direction=LT, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/lt" stack_frame_id=0} + %constant.1742.clone.1 = s32[]{:T(128)} constant(163840) + %closed_call.243 = s32[1,128,1]{1,2,0:T(1,128)} broadcast(%constant.1742.clone.1), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call" stack_frame_id=0} + %add.3593 = s32[1,128,1]{1,2,0:T(1,128)} add(%bitcast.1952, %closed_call.243), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/add" stack_frame_id=0} + %select_n.2276 = s32[1,128,1]{1,2,0:T(1,128)} select(%lt.836, %add.3593, %bitcast.1952), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/select_n" stack_frame_id=0} + ROOT %bitcast.1930 = s32[128]{0:T(128)S(1)} bitcast(%select_n.2276), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/select_n" stack_frame_id=0} +} + +%fused_computation.1130 (param_0.3318: s32[128]) -> s32[1024] { + %constant.4116 = s32[] constant(0), metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} + %broadcast.4264 = s32[1024]{0:T(1024)} broadcast(%constant.4116), dimensions={}, metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} + %param_0.3318 = s32[128]{0:T(128)S(1)} parameter(0) + %constant.4124 = s32[] constant(2147483647), metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} + %pad.329 = s32[1024]{0:T(1024)} pad(%param_0.3318, %constant.4124), padding=0_896, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/gather" stack_frame_id=0} + %constant.4109 = s32[] constant(163839), metadata={op_name="jit(train_step)/gather" stack_frame_id=0} + %broadcast.4256 = s32[1024]{0:T(1024)} broadcast(%constant.4109), dimensions={}, metadata={op_name="jit(train_step)/gather" stack_frame_id=0} + ROOT %clamp.57 = s32[1024]{0:T(1024)S(1)} clamp(%broadcast.4264, %pad.329, %broadcast.4256), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/gather" stack_frame_id=0} +} + +%fused_computation.10 (param_0.32: f32[163840,32], param_1.135: s32[1024]) -> f32[128,32] { + %param_0.32 = f32[163840,32]{1,0:T(8,128)} parameter(0) + %param_1.135 = s32[1024]{0:T(1024)S(1)} parameter(1) + %custom-call.32 = s32[1024]{0:T(1024)} custom-call(%param_1.135), custom_call_target="AssumeGatherIndicesInBound", operand_layout_constraints={s32[1024]{0:T(1024)}}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/gather" stack_frame_id=0} + %slice.1071 = s32[128]{0:T(128)} slice(%custom-call.32), slice={[0:128]}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/gather" stack_frame_id=0} + %reshape.4358 = s32[128]{0:T(128)} reshape(%slice.1071), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/select_n" stack_frame_id=0} + %transpose.915 = s32[128]{0:T(128)} transpose(%reshape.4358), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/select_n" stack_frame_id=0} + %gather.234 = f32[128,32]{1,0:T(8,128)} gather(%param_0.32, %transpose.915), offset_dims={1}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1,32}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/gather" stack_frame_id=0} + %transpose.914 = f32[128,32]{1,0:T(8,128)} transpose(%gather.234), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/gather" stack_frame_id=0} + ROOT %reshape.4357 = f32[128,32]{1,0:T(8,128)S(1)} reshape(%transpose.914), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/gather" stack_frame_id=0} +} + +%fused_computation.9 (param_0.29: f32[163840,32], param_1.133: s32[1024]) -> f32[128,32] { %param_0.29 = f32[163840,32]{1,0:T(8,128)} parameter(0) - %param_1.122 = s32[1024]{0:T(1024)S(1)} parameter(1) - %custom-call.27 = s32[1024]{0:T(1024)} custom-call(%param_1.122), custom_call_target="AssumeGatherIndicesInBound", operand_layout_constraints={s32[1024]{0:T(1024)}}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/gather" stack_frame_id=0} - %slice.906 = s32[512]{0:T(512)} slice(%custom-call.27), slice={[0:512]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/gather" stack_frame_id=0} - %reshape.3338 = s32[4,128]{1,0:T(4,128)} reshape(%slice.906), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/select_n" stack_frame_id=0} - %transpose.871 = s32[4,128]{1,0:T(4,128)} transpose(%reshape.3338), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/select_n" stack_frame_id=0} - %gather.191 = f32[4,128,32]{2,1,0:T(8,128)} gather(%param_0.29, %transpose.871), offset_dims={2}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=2, slice_sizes={1,32}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/gather" stack_frame_id=0} - %transpose.870 = f32[4,128,32]{2,1,0:T(8,128)} transpose(%gather.191), dimensions={0,1,2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/gather" stack_frame_id=0} - ROOT %reshape.3337 = f32[512,32]{1,0:T(8,128)S(1)} reshape(%transpose.870), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/gather" stack_frame_id=0} -} - -%fused_computation.10 (param_0.32: bf16[4096,512], param_1.126: s32[4096]) -> bf16[4096,512] { - %param_0.32 = bf16[4096,512]{1,0:T(8,128)(2,1)S(1)} parameter(0) - %param_1.126 = s32[4096]{0:T(1024)S(1)} parameter(1) - %custom-call.31 = s32[4096]{0:T(1024)} custom-call(%param_1.126), custom_call_target="AssumeGatherIndicesInBound", operand_layout_constraints={s32[4096]{0:T(1024)}}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} - %slice.910 = s32[4096]{0:T(1024)} slice(%custom-call.31), slice={[0:4096]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} - %reshape.3346 = s32[4096]{0:T(1024)} reshape(%slice.910), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/sort_activations/broadcast_in_dim" stack_frame_id=0} - %transpose.877 = s32[4096]{0:T(1024)} transpose(%reshape.3346), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/sort_activations/broadcast_in_dim" stack_frame_id=0} - %gather.193 = bf16[4096,512]{1,0:T(8,128)(2,1)} gather(%param_0.32, %transpose.877), offset_dims={1}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1,512}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} - %transpose.876 = bf16[4096,512]{1,0:T(8,128)(2,1)} transpose(%gather.193), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} - ROOT %reshape.3345 = bf16[4096,512]{1,0:T(8,128)(2,1)} reshape(%transpose.876), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} -} - -%fused_computation.11 (param_0.35: bf16[4096,512], param_1.128: s32[4096]) -> bf16[4096,512] { - %param_0.35 = bf16[4096,512]{1,0:T(8,128)(2,1)S(1)} parameter(0) - %param_1.128 = s32[4096]{0:T(1024)S(1)} parameter(1) - %custom-call.33 = s32[4096]{0:T(1024)} custom-call(%param_1.128), custom_call_target="AssumeGatherIndicesInBound", operand_layout_constraints={s32[4096]{0:T(1024)}}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} - %slice.912 = s32[4096]{0:T(1024)} slice(%custom-call.33), slice={[0:4096]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} - %reshape.3354 = s32[4096]{0:T(1024)} reshape(%slice.912), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/broadcast_in_dim" stack_frame_id=0} - %transpose.883 = s32[4096]{0:T(1024)} transpose(%reshape.3354), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/broadcast_in_dim" stack_frame_id=0} - %gather.195 = bf16[4096,512]{1,0:T(8,128)(2,1)} gather(%param_0.35, %transpose.883), offset_dims={1}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1,512}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} - %transpose.882 = bf16[4096,512]{1,0:T(8,128)(2,1)} transpose(%gather.195), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} - ROOT %reshape.3353 = bf16[4096,512]{1,0:T(8,128)(2,1)S(1)} reshape(%transpose.882), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} -} - -%fused_computation.12 (param_0.38: bf16[4096,512], param_1.130: s32[4096]) -> bf16[4096,512] { - %param_0.38 = bf16[4096,512]{1,0:T(8,128)(2,1)S(1)} parameter(0) - %param_1.130 = s32[4096]{0:T(1024)S(1)} parameter(1) - %custom-call.35 = s32[4096]{0:T(1024)} custom-call(%param_1.130), custom_call_target="AssumeGatherIndicesInBound", operand_layout_constraints={s32[4096]{0:T(1024)}}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} - %slice.914 = s32[4096]{0:T(1024)} slice(%custom-call.35), slice={[0:4096]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} - %reshape.3362 = s32[4096]{0:T(1024)} reshape(%slice.914), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/broadcast_in_dim" stack_frame_id=0} - %transpose.889 = s32[4096]{0:T(1024)} transpose(%reshape.3362), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/broadcast_in_dim" stack_frame_id=0} - %gather.197 = bf16[4096,512]{1,0:T(8,128)(2,1)} gather(%param_0.38, %transpose.889), offset_dims={1}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1,512}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} - %transpose.888 = bf16[4096,512]{1,0:T(8,128)(2,1)} transpose(%gather.197), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} - ROOT %reshape.3361 = bf16[4096,512]{1,0:T(8,128)(2,1)S(1)} reshape(%transpose.888), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} -} - -%fused_computation.13 (param_0.41: bf16[4096,512], param_1.132: s32[4096]) -> bf16[4096,512] { - %param_0.41 = bf16[4096,512]{1,0:T(8,128)(2,1)S(1)} parameter(0) - %param_1.132 = s32[4096]{0:T(1024)S(1)} parameter(1) - %custom-call.37 = s32[4096]{0:T(1024)} custom-call(%param_1.132), custom_call_target="AssumeGatherIndicesInBound", operand_layout_constraints={s32[4096]{0:T(1024)}}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} - %slice.916 = s32[4096]{0:T(1024)} slice(%custom-call.37), slice={[0:4096]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} - %reshape.3370 = s32[4096]{0:T(1024)} reshape(%slice.916), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/broadcast_in_dim" stack_frame_id=0} - %transpose.895 = s32[4096]{0:T(1024)} transpose(%reshape.3370), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/sort_activations/broadcast_in_dim" stack_frame_id=0} - %gather.199 = bf16[4096,512]{1,0:T(8,128)(2,1)} gather(%param_0.41, %transpose.895), offset_dims={1}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1,512}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} - %transpose.894 = bf16[4096,512]{1,0:T(8,128)(2,1)} transpose(%gather.199), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} - ROOT %reshape.3369 = bf16[4096,512]{1,0:T(8,128)(2,1)S(1)} reshape(%transpose.894), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} -} - -%fused_computation.15 (param_0.47: s32[256], param_1.124: s32[1024]) -> s32[263] { - %param_0.47 = s32[256]{0:T(256)S(1)} parameter(0) - %param_1.124 = s32[1024]{0:T(1024)S(1)} parameter(1) - %custom-call.29 = s32[1024]{0:T(1024)} custom-call(%param_1.124), custom_call_target="AssumeGatherIndicesInBound", operand_layout_constraints={s32[1024]{0:T(1024)}}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/jit(_take)/gather" stack_frame_id=0} - %slice.908 = s32[263]{0:T(512)} slice(%custom-call.29), slice={[0:263]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/jit(_take)/gather" stack_frame_id=0} - %reshape.3401 = s32[263]{0:T(512)} reshape(%slice.908), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/jit(_take)/broadcast_in_dim" stack_frame_id=0} - %transpose.911 = s32[263]{0:T(512)} transpose(%reshape.3401), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/jit(_take)/broadcast_in_dim" stack_frame_id=0} - %gather.204 = s32[263]{0:T(512)} gather(%param_0.47, %transpose.911), offset_dims={}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/jit(_take)/gather" stack_frame_id=0} - %transpose.910 = s32[263]{0:T(512)} transpose(%gather.204), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/jit(_take)/gather" stack_frame_id=0} - ROOT %reshape.3400 = s32[263]{0:T(512)S(1)} reshape(%transpose.910), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/jit(_take)/gather" stack_frame_id=0} -} - -%fused_computation.16 (param_0.50: s32[256], param_1.134: s32[1024]) -> s32[263] { - %param_0.50 = s32[256]{0:T(256)S(1)} parameter(0) - %param_1.134 = s32[1024]{0:T(1024)S(1)} parameter(1) - %custom-call.39 = s32[1024]{0:T(1024)} custom-call(%param_1.134), custom_call_target="AssumeGatherIndicesInBound", operand_layout_constraints={s32[1024]{0:T(1024)}}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/jit(_take)/gather" stack_frame_id=0} - %slice.918 = s32[263]{0:T(512)} slice(%custom-call.39), slice={[0:263]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/jit(_take)/gather" stack_frame_id=0} - %reshape.3424 = s32[263]{0:T(512)} reshape(%slice.918), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/jit(_take)/broadcast_in_dim" stack_frame_id=0} - %transpose.921 = s32[263]{0:T(512)} transpose(%reshape.3424), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/jit(_take)/broadcast_in_dim" stack_frame_id=0} - %gather.207 = s32[263]{0:T(512)} gather(%param_0.50, %transpose.921), offset_dims={}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/jit(_take)/gather" stack_frame_id=0} - %transpose.920 = s32[263]{0:T(512)} transpose(%gather.207), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/jit(_take)/gather" stack_frame_id=0} - ROOT %reshape.3423 = s32[263]{0:T(512)S(1)} reshape(%transpose.920), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/jit(_take)/gather" stack_frame_id=0} -} - -%region_173.198.clone (scatter-add.94: bf16[], scatter-add.96: bf16[]) -> bf16[] { - %scatter-add.94 = bf16[]{:T(256)} parameter(0), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/scatter-add"} - %scatter-add.96 = bf16[]{:T(256)} parameter(1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/scatter-add"} - ROOT %add.1918 = bf16[]{:T(256)} add(%scatter-add.94, %scatter-add.96), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/add" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.21 (param_0.55: bf16[129280,512], param_1.65: s32[512], param_2.24: bf16[512,512]) -> bf16[129280,512] { - %param_0.55 = bf16[129280,512]{1,0:T(8,128)(2,1)} parameter(0) - %param_1.65 = s32[512]{0:T(512)S(1)} parameter(1) - %reshape.3478 = s32[4,128]{1,0:T(4,128)} reshape(%param_1.65), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/broadcast_in_dim" stack_frame_id=0} - %transpose.954 = s32[4,128]{1,0:T(4,128)} transpose(%reshape.3478), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/broadcast_in_dim" stack_frame_id=0} - %param_2.24 = bf16[512,512]{1,0:T(8,128)(2,1)S(1)} parameter(2) - %reshape.3479 = bf16[4,128,512]{2,1,0:T(8,128)(2,1)} reshape(%param_2.24), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/dense_layers.wrapped_fn/while" stack_frame_id=0} - %transpose.955 = bf16[4,128,512]{2,1,0:T(8,128)(2,1)} transpose(%reshape.3479), dimensions={0,1,2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/dense_layers.wrapped_fn/while" stack_frame_id=0} - ROOT %scatter.73 = bf16[129280,512]{1,0:T(8,128)(2,1)} scatter(%param_0.55, %transpose.954, %transpose.955), update_window_dims={2}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=2, to_apply=%region_173.198.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/scatter-add" stack_frame_id=0} + %param_1.133 = s32[1024]{0:T(1024)S(1)} parameter(1) + %custom-call.30 = s32[1024]{0:T(1024)} custom-call(%param_1.133), custom_call_target="AssumeGatherIndicesInBound", operand_layout_constraints={s32[1024]{0:T(1024)}}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/gather" stack_frame_id=0} + %slice.1069 = s32[128]{0:T(128)} slice(%custom-call.30), slice={[0:128]}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/gather" stack_frame_id=0} + %reshape.4350 = s32[128]{0:T(128)} reshape(%slice.1069), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/select_n" stack_frame_id=0} + %transpose.909 = s32[128]{0:T(128)} transpose(%reshape.4350), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/select_n" stack_frame_id=0} + %gather.232 = f32[128,32]{1,0:T(8,128)} gather(%param_0.29, %transpose.909), offset_dims={1}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1,32}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/gather" stack_frame_id=0} + %transpose.908 = f32[128,32]{1,0:T(8,128)} transpose(%gather.232), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/gather" stack_frame_id=0} + ROOT %reshape.4349 = f32[128,32]{1,0:T(8,128)S(1)} reshape(%transpose.908), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/gather" stack_frame_id=0} +} + +%fused_computation.672 (param_0.2070: bf16[1,128,128,32,1], param_1.2330: bf16[1,128,128,32,1], param_2.4187: f32[128,32], param_3.2839: f32[128,32]) -> (bf16[1,128,128,32], bf16[1,128,128,32]) { + %param_1.2330 = bf16[1,128,128,32,1]{3,1,2,4,0:T(8,128)(2,1)} parameter(1) + %convert_element_type.2602 = f32[1,128,128,32,1]{3,1,2,4,0:T(8,128)} convert(%param_1.2330), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %bitcast.1420 = f32[1,128,128,32]{3,1,2,0:T(8,128)} bitcast(%convert_element_type.2602), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %param_0.2070 = bf16[1,128,128,32,1]{3,1,2,4,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.2601 = f32[1,128,128,32,1]{3,1,2,4,0:T(8,128)} convert(%param_0.2070), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %bitcast.1419 = f32[1,128,128,32]{3,1,2,0:T(8,128)} bitcast(%convert_element_type.2601), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %constant.1730.clone.87 = f32[]{:T(128)} constant(0) + %convert_element_type.2611 = f32[1,128,128,32]{3,1,2,0:T(8,128)} broadcast(%constant.1730.clone.87), dimensions={}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %mul.3926 = f32[1,128,128,32]{3,1,2,0:T(8,128)} multiply(%bitcast.1419, %convert_element_type.2611), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %add.3326 = f32[1,128,128,32]{3,1,2,0:T(8,128)} add(%bitcast.1420, %mul.3926), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/add" stack_frame_id=0} + %param_3.2839 = f32[128,32]{1,0:T(8,128)S(1)} parameter(3) + %broadcast_in_dim.1470 = f32[1,128,128,32]{3,1,2,0:T(8,128)} broadcast(%param_3.2839), dimensions={1,3}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/broadcast_in_dim" stack_frame_id=0} + %mul.3916 = f32[1,128,128,32]{3,1,2,0:T(8,128)} multiply(%add.3326, %broadcast_in_dim.1470), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %param_2.4187 = f32[128,32]{1,0:T(8,128)S(1)} parameter(2) + %broadcast_in_dim.1466 = f32[1,128,128,32]{3,1,2,0:T(8,128)} broadcast(%param_2.4187), dimensions={1,3}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/broadcast_in_dim" stack_frame_id=0} + %mul.3851 = f32[1,128,128,32]{3,1,2,0:T(8,128)} multiply(%bitcast.1419, %broadcast_in_dim.1466), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %mul.3850 = f32[1,128,128,32]{3,1,2,0:T(8,128)} subtract(%mul.3916, %mul.3851), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %convert.527 = bf16[1,128,128,32]{3,1,2,0:T(8,128)(2,1)} convert(%mul.3850), metadata={op_name="convert.378" stack_frame_id=0} + %mul.3915.clone.1 = f32[1,128,128,32]{3,1,2,0:T(8,128)} multiply(%add.3326, %broadcast_in_dim.1466), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %mul.3858.clone.1 = f32[1,128,128,32]{3,1,2,0:T(8,128)} multiply(%bitcast.1419, %broadcast_in_dim.1470), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %mul.3857.clone.1 = f32[1,128,128,32]{3,1,2,0:T(8,128)} add(%mul.3915.clone.1, %mul.3858.clone.1), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %convert.528.clone.1 = bf16[1,128,128,32]{3,1,2,0:T(8,128)(2,1)} convert(%mul.3857.clone.1), metadata={op_name="convert.379" stack_frame_id=0} + ROOT %tuple.667 = (bf16[1,128,128,32]{3,1,2,0:T(8,128)(2,1)}, bf16[1,128,128,32]{3,1,2,0:T(8,128)(2,1)}) tuple(%convert.527, %convert.528.clone.1) +} + +%fused_computation.657 (param_0.2067: bf16[1,128,128,32], param_1.4743: bf16[1,128,128,32], param_2.3579: bf16[1,128,128,128]) -> bf16[128,128,192] { + %param_2.3579 = bf16[1,128,128,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(2) + %constant.1746.clone.23 = bf16[]{:T(256)} constant(-inf) + %pad.288 = bf16[1,128,128,192]{3,1,2,0:T(8,128)(2,1)} pad(%param_2.3579, %constant.1746.clone.23), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/concatenate" stack_frame_id=0} + %convert.684 = f32[1,128,128,192]{3,1,2,0:T(8,128)} convert(%pad.288) + %param_1.4743 = bf16[1,128,128,32]{3,1,2,0:T(8,128)(2,1)} parameter(1) + %pad.287 = bf16[1,128,128,192]{3,1,2,0:T(8,128)(2,1)} pad(%param_1.4743, %constant.1746.clone.23), padding=0_0x0_0x0_0x128_32, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/concatenate" stack_frame_id=0} + %convert.685 = f32[1,128,128,192]{3,1,2,0:T(8,128)} convert(%pad.287) + %maximum.53 = f32[1,128,128,192]{3,1,2,0:T(8,128)} maximum(%convert.684, %convert.685), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/concatenate" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","128","192"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %param_0.2067 = bf16[1,128,128,32]{3,1,2,0:T(8,128)(2,1)} parameter(0) + %pad.286 = bf16[1,128,128,192]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.2067, %constant.1746.clone.23), padding=0_0x0_0x0_0x160_0, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/concatenate" stack_frame_id=0} + %convert.686 = f32[1,128,128,192]{3,1,2,0:T(8,128)} convert(%pad.286) + %maximum.52 = f32[1,128,128,192]{3,1,2,0:T(8,128)} maximum(%maximum.53, %convert.686), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/concatenate" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","128","192"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %constant.1684.clone.2 = bf16[]{:T(256)} constant(0.1348), metadata={stack_frame_id=0} + %mul.3949 = bf16[1,128,128,192]{3,1,2,0:T(8,128)(2,1)} broadcast(%constant.1684.clone.2), dimensions={}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %convert.687 = f32[1,128,128,192]{3,1,2,0:T(8,128)} convert(%mul.3949) + %mul.3833 = f32[1,128,128,192]{3,1,2,0:T(8,128)} multiply(%maximum.52, %convert.687), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","128","192"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %convert.688 = bf16[1,128,128,192]{3,1,2,0:T(8,128)(2,1)} convert(%mul.3833) + ROOT %bitcast.1354 = bf16[128,128,192]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%convert.688), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/shard_map/vmap(static_sparsity_splash)/jit(_splash_attention)/splash_mha_fwd_segmented_residuals/splash_mha_fwd_segmented_residuals/squeeze" stack_frame_id=0} +} + +%fused_computation.973 (param_0.2787: bf16[1,128,512]) -> bf16[128,512] { + %param_0.2787 = bf16[1,128,512]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.1737 = bf16[128,512]{1,0:T(8,128)(2,1)} bitcast(%param_0.2787), metadata={stack_frame_id=0} +} + +%fused_computation.778.clone.clone (param_0.2260: bf16[512,1,576]) -> bf16[512,576] { + %param_0.2260 = bf16[512,1,576]{0,2,1:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.1529 = bf16[512,576]{0,1:T(8,128)(2,1)} bitcast(%param_0.2260), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} +} + +%fused_computation.783 (param_0.2786: bf16[512,1,576], param_1.3050: bf16[1,128,512]) -> bf16[1,128,576] { + %param_1.3050 = bf16[1,128,512]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %fusion.690 = bf16[128,512]{1,0:T(8,128)(2,1)} fusion(%param_1.3050), kind=kLoop, calls=%fused_computation.973, metadata={stack_frame_id=0} + %param_0.2786 = bf16[512,1,576]{0,2,1:T(8,128)(2,1)S(1)} parameter(0) + %fusion.515 = bf16[512,576]{0,1:T(8,128)(2,1)} fusion(%param_0.2786), kind=kLoop, calls=%fused_computation.778.clone.clone, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} + %convolution.209 = bf16[128,576]{0,1:T(8,128)(2,1)} convolution(%fusion.690, %fusion.515), dim_labels=bf_io->bf, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/dot_general" stack_frame_id=0} + ROOT %bitcast.1527 = bf16[1,128,576]{1,2,0:T(8,128)(2,1)S(1)} bitcast(%convolution.209), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/dot_general" stack_frame_id=0} +} + +%fused_computation.1429 (param_0.4320: bf16[1,128,1,32,2]) -> (bf16[1,128,1,32,1], bf16[1,128,1,32,1]) { + %param_0.4320 = bf16[1,128,1,32,2]{1,4,3,2,0:T(2,128)(2,1)S(1)} parameter(0) + %slice.1339 = bf16[1,128,1,32,1]{1,4,3,2,0:T(2,128)(2,1)S(1)} slice(%param_0.4320), slice={[0:1], [0:128], [0:1], [0:32], [0:1]}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/slice" stack_frame_id=0} + %slice.1340 = bf16[1,128,1,32,1]{1,4,3,2,0:T(2,128)(2,1)S(1)} slice(%param_0.4320), slice={[0:1], [0:128], [0:1], [0:32], [1:2]}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/slice" stack_frame_id=0} + ROOT %tuple.696 = (bf16[1,128,1,32,1]{1,4,3,2,0:T(2,128)(2,1)S(1)}, bf16[1,128,1,32,1]{1,4,3,2,0:T(2,128)(2,1)S(1)}) tuple(%slice.1339, %slice.1340) +} + +%fused_computation.1027 (param_0.2945: f32[128,32], param_1.3280: f32[128,32], param_2.2308: bf16[1,128,1,32,1], param_3.1432: bf16[1,128,1,32,1]) -> (bf16[1,128,32], bf16[1,128,32]) { + %param_3.1432 = bf16[1,128,1,32,1]{3,1,4,2,0:T(8,128)(2,1)S(1)} parameter(3) + %convert_element_type.2801 = f32[1,128,1,32,1]{3,1,4,2,0:T(8,128)} convert(%param_3.1432), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %bitcast.1800 = f32[1,128,32]{2,1,0:T(8,128)} bitcast(%convert_element_type.2801), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %param_2.2308 = bf16[1,128,1,32,1]{3,1,4,2,0:T(8,128)(2,1)S(1)} parameter(2) + %convert_element_type.2799 = f32[1,128,1,32,1]{3,1,4,2,0:T(8,128)} convert(%param_2.2308), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %bitcast.1798 = f32[1,128,32]{2,1,0:T(8,128)} bitcast(%convert_element_type.2799), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %constant.1730.clone.86 = f32[]{:T(128)} constant(0) + %convert_element_type.2819 = f32[1,128,32]{2,1,0:T(8,128)} broadcast(%constant.1730.clone.86), dimensions={}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %mul.4320 = f32[1,128,32]{2,1,0:T(8,128)} multiply(%bitcast.1798, %convert_element_type.2819), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %add.3426 = f32[1,128,32]{2,1,0:T(8,128)} add(%bitcast.1800, %mul.4320), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/add" stack_frame_id=0} + %param_1.3280 = f32[128,32]{1,0:T(8,128)S(1)} parameter(1) + %bitcast.1789 = f32[1,128,32]{2,1,0:T(8,128)} bitcast(%param_1.3280), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/gather" stack_frame_id=0} + %mul.4313 = f32[1,128,32]{2,1,0:T(8,128)} multiply(%add.3426, %bitcast.1789), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %param_0.2945 = f32[128,32]{1,0:T(8,128)S(1)} parameter(0) + %bitcast.1783 = f32[1,128,32]{2,1,0:T(8,128)} bitcast(%param_0.2945), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/gather" stack_frame_id=0} + %mul.4307 = f32[1,128,32]{2,1,0:T(8,128)} multiply(%bitcast.1798, %bitcast.1783), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %mul.4306 = f32[1,128,32]{2,1,0:T(8,128)} subtract(%mul.4313, %mul.4307), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %convert.536 = bf16[1,128,32]{2,1,0:T(8,128)(2,1)S(1)} convert(%mul.4306), metadata={op_name="convert.376" stack_frame_id=0} + %mul.4300.clone.1 = f32[1,128,32]{2,1,0:T(8,128)} multiply(%add.3426, %bitcast.1783), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %mul.4311.clone.1 = f32[1,128,32]{2,1,0:T(8,128)} multiply(%bitcast.1798, %bitcast.1789), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %mul.4299.clone.1 = f32[1,128,32]{2,1,0:T(8,128)} add(%mul.4300.clone.1, %mul.4311.clone.1), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %convert.535.clone.1 = bf16[1,128,32]{2,1,0:T(8,128)(2,1)S(1)} convert(%mul.4299.clone.1), metadata={op_name="convert.377" stack_frame_id=0} + ROOT %tuple.688 = (bf16[1,128,32]{2,1,0:T(8,128)(2,1)S(1)}, bf16[1,128,32]{2,1,0:T(8,128)(2,1)S(1)}) tuple(%convert.536, %convert.535.clone.1) +} + +%fused_computation.1065 (param_0.3062: bf16[1,128,32], param_1.4733: bf16[1,128,32]) -> bf16[128,64] { + %param_1.4733 = bf16[1,128,32]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %constant.1746.clone.5 = bf16[]{:T(256)} constant(-inf) + %pad.316 = bf16[1,128,64]{2,1,0:T(8,128)(2,1)} pad(%param_1.4733, %constant.1746.clone.5), padding=0_0x0_0x0_32, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/concatenate" stack_frame_id=0} + %convert.798 = f32[1,128,64]{2,1,0:T(8,128)} convert(%pad.316) + %param_0.3062 = bf16[1,128,32]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %pad.315 = bf16[1,128,64]{2,1,0:T(8,128)(2,1)} pad(%param_0.3062, %constant.1746.clone.5), padding=0_0x0_0x32_0, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/concatenate" stack_frame_id=0} + %convert.799 = f32[1,128,64]{2,1,0:T(8,128)} convert(%pad.315) + %maximum.65 = f32[1,128,64]{2,1,0:T(8,128)} maximum(%convert.798, %convert.799), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/concatenate" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","64"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %convert.800 = bf16[1,128,64]{2,1,0:T(8,128)(2,1)} convert(%maximum.65) + ROOT %bitcast.1840 = bf16[128,64]{1,0:T(8,128)(2,1)S(1)} bitcast(%convert.800), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/concatenate" stack_frame_id=0} +} + +%region_44.53 (reduce_sum.110: f32[], reduce_sum.111: f32[]) -> f32[] { + %reduce_sum.110 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/reduce_sum"} + %reduce_sum.111 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/reduce_sum"} + ROOT %reduce_sum.115 = f32[]{:T(128)} add(%reduce_sum.110, %reduce_sum.111), metadata={op_name="checkpoint/rematted_computation/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%fused_computation.892 (param_0.4489: bf16[1,128,576]) -> f32[128] { + %param_0.4489 = bf16[1,128,576]{1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %slice.1225 = bf16[1,128,512]{1,2,0:T(8,128)(2,1)} slice(%param_0.4489), slice={[0:1], [0:128], [0:512]}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/split" stack_frame_id=0} + %convert_element_type.2696 = f32[1,128,512]{1,2,0:T(8,128)} convert(%slice.1225), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %square.626 = f32[1,128,512]{1,2,0:T(8,128)} multiply(%convert_element_type.2696, %convert_element_type.2696), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/square" stack_frame_id=0} + %constant.1730.clone.85 = f32[]{:T(128)} constant(0) + ROOT %reduce.763 = f32[128]{0:T(128)S(1)} reduce(%square.626, %constant.1730.clone.85), dimensions={0,2}, to_apply=%region_44.53, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/reduce_sum" stack_frame_id=0} +} + +%fused_computation.1330 (param_0.3991: f32[128]) -> f32[128] { + %param_0.3991 = f32[128]{0:T(128)S(1)} parameter(0) + %constant.1744.clone.7 = f32[]{:T(128)} constant(0.001953125) + %broadcast.4398 = f32[128]{0:T(128)} broadcast(%constant.1744.clone.7), dimensions={}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call" stack_frame_id=0} + %div.2465 = f32[128]{0:T(128)} multiply(%param_0.3991, %broadcast.4398), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/div" stack_frame_id=0} + %constant.1745.clone.9 = f32[]{:T(128)} constant(1e-06) + %broadcast.4387 = f32[128]{0:T(128)} broadcast(%constant.1745.clone.9), dimensions={}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call" stack_frame_id=0} + %add.3644 = f32[128]{0:T(128)} add(%div.2465, %broadcast.4387), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/add" stack_frame_id=0} + %bitcast.1951 = f32[1,128]{1,0:T(1,128)} bitcast(%add.3644), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/add" stack_frame_id=0} + %rsqrt.305 = f32[1,128]{1,0:T(1,128)} rsqrt(%bitcast.1951), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/rsqrt" stack_frame_id=0} + ROOT %bitcast.1923 = f32[128]{0:T(128)S(1)} bitcast(%rsqrt.305), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/rsqrt" stack_frame_id=0} +} + +%fused_computation.958 (param_0.3001: f32[128], param_1.3336: bf16[1,128,576], param_2.2341: bf16[512]) -> bf16[128,512,1] { + %param_2.2341 = bf16[512]{0:T(512)(128)(2,1)S(1)} parameter(2) + %dot_general.791 = bf16[1,128,512]{1,2,0:T(8,128)(2,1)} broadcast(%param_2.2341), dimensions={2}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/...k,k->...k/dot_general" stack_frame_id=0} + %convert.831 = f32[1,128,512]{1,2,0:T(8,128)} convert(%dot_general.791) + %param_1.3336 = bf16[1,128,576]{1,2,0:T(8,128)(2,1)S(1)} parameter(1) + %slice.1222 = bf16[1,128,512]{1,2,0:T(8,128)(2,1)} slice(%param_1.3336), slice={[0:1], [0:128], [0:512]}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/split" stack_frame_id=0} + %convert_element_type.2777 = f32[1,128,512]{1,2,0:T(8,128)} convert(%slice.1222), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %param_0.3001 = f32[128]{0:T(128)S(1)} parameter(0) + %mul.4246 = f32[1,128,512]{1,2,0:T(8,128)} broadcast(%param_0.3001), dimensions={1}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %mul.4245 = f32[1,128,512]{1,2,0:T(8,128)} multiply(%convert_element_type.2777, %mul.4246), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %convert_element_type.2776 = bf16[1,128,512]{1,2,0:T(8,128)(2,1)} convert(%mul.4245), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %convert.832 = f32[1,128,512]{1,2,0:T(8,128)} convert(%convert_element_type.2776) + %dot_general.768 = f32[1,128,512]{1,2,0:T(8,128)} multiply(%convert.831, %convert.832), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/...k,k->...k/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","512"],"tuple_shapes":[],"layout":{"minor_to_major":["1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %convert.833 = bf16[1,128,512]{1,2,0:T(8,128)(2,1)} convert(%dot_general.768) + ROOT %bitcast.1714 = bf16[128,512,1]{0,1,2:T(8,128)(2,1)} bitcast(%convert.833) +} + +%fused_computation.592.clone.clone (param_0.2036: bf16[512,1,128,256]) -> bf16[512,128,256] { + %param_0.2036 = bf16[512,1,128,256]{3,0,2,1:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.1434 = bf16[512,128,256]{2,0,1:T(8,128)(2,1)} bitcast(%param_0.2036), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} +} + +%fused_computation.698 (param_0.2759: bf16[512,1,128,256], param_1.3335: f32[128], param_2.2340: bf16[1,128,576], param_3.1371: bf16[512]) -> bf16[1,128,128,256] { + %param_1.3335 = f32[128]{0:T(128)S(1)} parameter(1) + %param_2.2340 = bf16[1,128,576]{1,2,0:T(8,128)(2,1)S(1)} parameter(2) + %param_3.1371 = bf16[512]{0:T(512)(128)(2,1)S(1)} parameter(3) + %fusion.677 = bf16[128,512,1]{0,1,2:T(8,128)(2,1)} fusion(%param_1.3335, %param_2.2340, %param_3.1371), kind=kLoop, calls=%fused_computation.958 + %param_0.2759 = bf16[512,1,128,256]{3,0,2,1:T(8,128)(2,1)S(1)} parameter(0) + %fusion.458 = bf16[512,128,256]{2,0,1:T(8,128)(2,1)} fusion(%param_0.2759), kind=kLoop, calls=%fused_computation.592.clone.clone, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} + %convolution.181 = bf16[128,128,256]{2,0,1:T(8,128)(2,1)} convolution(%fusion.677, %fusion.458), window={size=128 pad=127_127 rhs_reversal=1}, dim_labels=bf0_i0o->b0f, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/dot_general" stack_frame_id=0} + ROOT %bitcast.1432 = bf16[1,128,128,256]{3,1,2,0:T(8,128)(2,1)} bitcast(%convolution.181), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/dot_general" stack_frame_id=0} +} + +%fused_computation.705 (param_0.2057: bf16[1,128,128,256]) -> (bf16[128,128,128], bf16[1,128,128,128]) { + %param_0.2057 = bf16[1,128,128,256]{3,1,2,0:T(8,128)(2,1)} parameter(0) + %slice.1207 = bf16[1,128,128,128]{3,1,2,0:T(8,128)(2,1)} slice(%param_0.2057), slice={[0:1], [0:128], [0:128], [128:256]}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/split" stack_frame_id=0} + %bitcast.1439 = bf16[128,128,128]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%slice.1207), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/shard_map/vmap(static_sparsity_splash)/jit(_splash_attention)/splash_mha_fwd_segmented_residuals/splash_mha_fwd_segmented_residuals/squeeze" stack_frame_id=0} + %slice.1333 = bf16[1,128,128,128]{3,1,2,0:T(8,128)(2,1)S(1)} slice(%param_0.2057), slice={[0:1], [0:128], [0:128], [0:128]}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/split" stack_frame_id=0} + ROOT %tuple.672 = (bf16[128,128,128]{2,1,0:T(8,128)(2,1)S(1)}, bf16[1,128,128,128]{3,1,2,0:T(8,128)(2,1)S(1)}) tuple(%bitcast.1439, %slice.1333) +} + +%fused_computation.658 (param_0.1913: bf16[1,128,128,64], param_1.4735: bf16[1,128,128,128]) -> bf16[128,128,192] { + %param_1.4735 = bf16[1,128,128,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) + %constant.1746.clone.7 = bf16[]{:T(256)} constant(-inf) + %pad.290 = bf16[1,128,128,192]{3,1,2,0:T(8,128)(2,1)} pad(%param_1.4735, %constant.1746.clone.7), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/concatenate" stack_frame_id=0} + %convert.689 = f32[1,128,128,192]{3,1,2,0:T(8,128)} convert(%pad.290) + %param_0.1913 = bf16[1,128,128,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %pad.289 = bf16[1,128,128,192]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1913, %constant.1746.clone.7), padding=0_0x0_0x0_0x128_0, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/concatenate" stack_frame_id=0} + %convert.690 = f32[1,128,128,192]{3,1,2,0:T(8,128)} convert(%pad.289) + %maximum.54 = f32[1,128,128,192]{3,1,2,0:T(8,128)} maximum(%convert.689, %convert.690), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/concatenate" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","128","192"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %convert.691 = bf16[1,128,128,192]{3,1,2,0:T(8,128)(2,1)} convert(%maximum.54) + ROOT %bitcast.1355 = bf16[128,128,192]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%convert.691), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/shard_map/vmap(static_sparsity_splash)/jit(_splash_attention)/splash_mha_fwd_segmented_residuals/splash_mha_fwd_segmented_residuals/squeeze" stack_frame_id=0} +} + +%fused_computation.710.clone.clone.clone.clone.clone.clone.clone.clone (param_0.4337: bf16[128,128,128]) -> bf16[128,128,128] { + %param_0.4337 = bf16[128,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.2105 = bf16[128,128,128]{2,0,1:T(8,128)(2,1)} bitcast(%param_0.4337) +} + +%fused_computation.640.clone.clone.clone.clone.clone.clone.clone.clone (param_0.4336: bf16[128,1,128,512]) -> bf16[128,128,512] { + %param_0.4336 = bf16[128,1,128,512]{3,2,0,1:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.2104 = bf16[128,128,512]{2,1,0:T(8,128)(2,1)} bitcast(%param_0.4336), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} +} + +%region_45.56 (reduce_sum.116: f32[], reduce_sum.117: f32[]) -> f32[] { + %reduce_sum.116 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/reduce_sum"} + %reduce_sum.117 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/reduce_sum"} + ROOT %reduce_sum.118 = f32[]{:T(128)} add(%reduce_sum.116, %reduce_sum.117), metadata={op_name="checkpoint/rematted_computation/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%fused_computation.917 (param_0.4488: bf16[1,128,512], param_1.5220: bf16[128,1,128,512], param_2.4186: bf16[128,128,128]) -> (f32[128], f32[1,128,512]) { + %param_0.4488 = bf16[1,128,512]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert.734 = f32[1,128,512]{2,1,0:T(8,128)} convert(%param_0.4488) + %param_2.4186 = bf16[128,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %fusion.600.clone.1 = bf16[128,128,128]{2,0,1:T(8,128)(2,1)} fusion(%param_2.4186), kind=kLoop, calls=%fused_computation.710.clone.clone.clone.clone.clone.clone.clone.clone + %param_1.5220 = bf16[128,1,128,512]{3,2,0,1:T(8,128)(2,1)S(1)} parameter(1) + %fusion.599.clone.1 = bf16[128,128,512]{2,1,0:T(8,128)(2,1)} fusion(%param_1.5220), kind=kLoop, calls=%fused_computation.640.clone.clone.clone.clone.clone.clone.clone.clone, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} + %convolution.261.clone.1 = bf16[128,512,1]{1,0,2:T(8,128)(2,1)} convolution(%fusion.600.clone.1, %fusion.599.clone.1), window={size=128}, dim_labels=b0f_0io->bf0, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/dot_general" stack_frame_id=0} + %bitcast.1620.clone.1 = bf16[1,128,512]{2,1,0:T(8,128)(2,1)} bitcast(%convolution.261.clone.1), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/dot_general" stack_frame_id=0} + %convert.735 = f32[1,128,512]{2,1,0:T(8,128)} convert(%bitcast.1620.clone.1) + %add.3370.clone.1 = f32[1,128,512]{2,1,0:T(8,128)} add(%convert.734, %convert.735), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/add" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","512"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %convert_element_type.2639.clone.1 = f32[1,128,512]{2,1,0:T(8,128)S(1)} convert(%add.3370.clone.1), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %square.638 = f32[1,128,512]{2,1,0:T(8,128)} multiply(%convert_element_type.2639.clone.1, %convert_element_type.2639.clone.1), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/square" stack_frame_id=0} + %constant.1730.clone.84 = f32[]{:T(128)} constant(0) + %reduce.770 = f32[128]{0:T(128)S(1)} reduce(%square.638, %constant.1730.clone.84), dimensions={0,2}, to_apply=%region_45.56, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/reduce_sum" stack_frame_id=0} + ROOT %tuple.724 = (f32[128]{0:T(128)S(1)}, f32[1,128,512]{2,1,0:T(8,128)S(1)}) tuple(%reduce.770, %convert_element_type.2639.clone.1) +} + +%fused_computation.1328 (param_0.3993: f32[128]) -> f32[128] { + %param_0.3993 = f32[128]{0:T(128)S(1)} parameter(0) + %constant.1744.clone.6 = f32[]{:T(128)} constant(0.001953125) + %broadcast.4399 = f32[128]{0:T(128)} broadcast(%constant.1744.clone.6), dimensions={}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call" stack_frame_id=0} + %div.2469 = f32[128]{0:T(128)} multiply(%param_0.3993, %broadcast.4399), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/div" stack_frame_id=0} + %constant.1745.clone.7 = f32[]{:T(128)} constant(1e-06) + %broadcast.4389 = f32[128]{0:T(128)} broadcast(%constant.1745.clone.7), dimensions={}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call" stack_frame_id=0} + %add.3648 = f32[128]{0:T(128)} add(%div.2469, %broadcast.4389), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/add" stack_frame_id=0} + %bitcast.1949 = f32[1,128]{1,0:T(1,128)} bitcast(%add.3648), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/add" stack_frame_id=0} + %rsqrt.303 = f32[1,128]{1,0:T(1,128)} rsqrt(%bitcast.1949), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/rsqrt" stack_frame_id=0} + ROOT %bitcast.1922 = f32[128]{0:T(128)S(1)} bitcast(%rsqrt.303), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/rsqrt" stack_frame_id=0} +} + +%fused_computation.995 (param_0.3102: bf16[512], param_1.3443: f32[1,128,512], param_2.2438: f32[128], param_3.2756: f32[128], param_4.2113: bf16[1,128,512]) -> (bf16[128,512], bf16[128,512]) { + %param_0.3102 = bf16[512]{0:T(512)(128)(2,1)S(1)} parameter(0) + %dot_general.809 = bf16[1,128,512]{2,1,0:T(8,128)(2,1)} broadcast(%param_0.3102), dimensions={2}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/...k,k->...k/dot_general" stack_frame_id=0} + %convert.767 = f32[1,128,512]{2,1,0:T(8,128)} convert(%dot_general.809) + %param_1.3443 = f32[1,128,512]{2,1,0:T(8,128)S(1)} parameter(1) + %param_2.2438 = f32[128]{0:T(128)S(1)} parameter(2) + %mul.4395 = f32[1,128,512]{2,1,0:T(8,128)} broadcast(%param_2.2438), dimensions={1}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %mul.4394 = f32[1,128,512]{2,1,0:T(8,128)} multiply(%param_1.3443, %mul.4395), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %convert_element_type.2826 = bf16[1,128,512]{2,1,0:T(8,128)(2,1)} convert(%mul.4394), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %convert.768 = f32[1,128,512]{2,1,0:T(8,128)} convert(%convert_element_type.2826) + %dot_general.754 = f32[1,128,512]{2,1,0:T(8,128)} multiply(%convert.767, %convert.768), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/...k,k->...k/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","512"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %convert.769 = bf16[1,128,512]{2,1,0:T(8,128)(2,1)} convert(%dot_general.754) + %bitcast.1744 = bf16[128,512]{1,0:T(8,128)(2,1)S(1)} bitcast(%convert.769), metadata={stack_frame_id=0} + %param_4.2113 = bf16[1,128,512]{2,1,0:T(8,128)(2,1)S(1)} parameter(4) + %convert_element_type.2739.clone.1 = f32[1,128,512]{2,1,0:T(8,128)} convert(%param_4.2113), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %param_3.2756 = f32[128]{0:T(128)S(1)} parameter(3) + %mul.4198.clone.1 = f32[1,128,512]{2,1,0:T(8,128)} broadcast(%param_3.2756), dimensions={1}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %mul.4174.clone.1 = f32[1,128,512]{2,1,0:T(8,128)} multiply(%convert_element_type.2739.clone.1, %mul.4198.clone.1), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %convert_element_type.2738.clone.1 = bf16[1,128,512]{2,1,0:T(8,128)(2,1)} convert(%mul.4174.clone.1), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %convert.770 = f32[1,128,512]{2,1,0:T(8,128)} convert(%convert_element_type.2738.clone.1) + %dot_general.749.clone.1 = f32[1,128,512]{2,1,0:T(8,128)} multiply(%convert.767, %convert.770), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/...k,k->...k/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","512"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %convert.771 = bf16[1,128,512]{2,1,0:T(8,128)(2,1)} convert(%dot_general.749.clone.1) + %bitcast.1706.clone.1 = bf16[128,512]{1,0:T(8,128)(2,1)S(1)} bitcast(%convert.771), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/reshape" stack_frame_id=0} + ROOT %tuple.707 = (bf16[128,512]{1,0:T(8,128)(2,1)S(1)}, bf16[128,512]{1,0:T(8,128)(2,1)S(1)}) tuple(%bitcast.1744, %bitcast.1706.clone.1) +} + +%bitcast_fusion.5 (bitcast_input.5: bf16[128,512]) -> bf16[128,512] { + %bitcast_input.5 = bf16[128,512]{1,0:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.2163 = bf16[128,512]{1,0:T(8,128)(2,1)} bitcast(%bitcast_input.5) +} + +%fused_computation.906.clone.clone (param_0.4348: bf16[512,1,256]) -> bf16[512,256] { + %param_0.4348 = bf16[512,1,256]{2,0,1:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.2115 = bf16[512,256]{1,0:T(8,128)(2,1)} bitcast(%param_0.4348), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} +} + +%fused_computation.1021 (param_0.4347: bf16[256], param_1.5120: bf16[128,256], param_2.4120: bf16[128,512], param_3.2779: bf16[512,1,256]) -> (bf16[1,128,256], bf16[1,128,256], bf16[128,256]) { + %constant.1687.clone.6.clone.1 = bf16[]{:T(256)} constant(1), metadata={stack_frame_id=0} + %broadcast.4227.clone.1 = bf16[128,256]{1,0:T(8,128)(2,1)} broadcast(%constant.1687.clone.6.clone.1), dimensions={}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call" stack_frame_id=0} + %convert.779 = f32[128,256]{1,0:T(8,128)} convert(%broadcast.4227.clone.1) + %param_2.4120 = bf16[128,512]{1,0:T(8,128)(2,1)S(1)} parameter(2) + %fusion.896 = bf16[128,512]{1,0:T(8,128)(2,1)} fusion(%param_2.4120), kind=kLoop, calls=%bitcast_fusion.5 + %param_3.2779 = bf16[512,1,256]{2,0,1:T(8,128)(2,1)S(1)} parameter(3) + %fusion.650.clone.1 = bf16[512,256]{1,0:T(8,128)(2,1)} fusion(%param_3.2779), kind=kLoop, calls=%fused_computation.906.clone.clone, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} + %convolution.270.clone.1 = bf16[128,256]{1,0:T(8,128)(2,1)} convolution(%fusion.896, %fusion.650.clone.1), dim_labels=bf_io->bf, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/dot_general" stack_frame_id=0} + %convert.778 = f32[128,256]{1,0:T(8,128)} convert(%convolution.270.clone.1) + %neg.211.clone.1 = f32[128,256]{1,0:T(8,128)} negate(%convert.778), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/neg" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["128","256"],"tuple_shapes":[],"layout":{"minor_to_major":["1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false]}}} + %exp.506.clone.1 = f32[128,256]{1,0:T(8,128)} exponential(%neg.211.clone.1), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/exp" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["128","256"],"tuple_shapes":[],"layout":{"minor_to_major":["1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false]}}} + %add.3374.clone.1 = f32[128,256]{1,0:T(8,128)} add(%exp.506.clone.1, %convert.779), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/add" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["128","256"],"tuple_shapes":[],"layout":{"minor_to_major":["1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false]}}} + %div.2161.clone.1 = f32[128,256]{1,0:T(8,128)} divide(%convert.779, %add.3374.clone.1), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/div" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["128","256"],"tuple_shapes":[],"layout":{"minor_to_major":["1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false]}}} + %convert.780 = bf16[128,256]{1,0:T(8,128)(2,1)} convert(%div.2161.clone.1) + %bitcast.1773 = bf16[1,128,256]{2,1,0:T(8,128)(2,1)} bitcast(%convert.780), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/div" stack_frame_id=0} + %convert.781 = f32[1,128,256]{2,1,0:T(8,128)} convert(%bitcast.1773) + %param_0.4347 = bf16[256]{0:T(256)(128)(2,1)S(1)} parameter(0) + %add.3446 = bf16[1,128,256]{2,1,0:T(8,128)(2,1)} broadcast(%param_0.4347), dimensions={2}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/add" stack_frame_id=0} + %convert.782 = f32[1,128,256]{2,1,0:T(8,128)} convert(%add.3446) + %add.3421 = f32[1,128,256]{2,1,0:T(8,128)} add(%convert.781, %convert.782), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/add" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","256"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %convert.784 = bf16[1,128,256]{2,1,0:T(8,128)(2,1)} convert(%add.3421) + %param_1.5120 = bf16[128,256]{1,0:T(8,128)(2,1)S(1)} parameter(1) + %bitcast.1802.clone.1 = bf16[1,128,256]{2,1,0:T(8,128)(2,1)} bitcast(%param_1.5120), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/div" stack_frame_id=0} + %convert.783 = f32[1,128,256]{2,1,0:T(8,128)} convert(%bitcast.1802.clone.1) + %add.3427.clone.1 = f32[1,128,256]{2,1,0:T(8,128)} add(%convert.783, %convert.782), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/add" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","256"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %convert.785 = bf16[1,128,256]{2,1,0:T(8,128)(2,1)S(1)} convert(%add.3427.clone.1) + ROOT %tuple.727 = (bf16[1,128,256]{2,1,0:T(8,128)(2,1)}, bf16[1,128,256]{2,1,0:T(8,128)(2,1)S(1)}, bf16[128,256]{1,0:T(8,128)(2,1)}) tuple(%convert.784, %convert.785, %convert.780) } %region_12.18 (top_k.0: bf16[], top_k.6: bf16[], top_k.7: s32[], top_k.8: s32[]) -> pred[] { - %constant.1387 = s32[]{:T(128)} constant(0) - %constant.1388 = s32[]{:T(128)} constant(2147483647) + %constant.1521 = s32[]{:T(128)} constant(0) + %constant.1522 = s32[]{:T(128)} constant(2147483647) %top_k.0 = bf16[]{:T(256)} parameter(0), metadata={op_name="top_k"} %top_k.6 = bf16[]{:T(256)} parameter(1), metadata={op_name="top_k"} %top_k.7 = s32[]{:T(128)} parameter(2), metadata={op_name="top_k"} %top_k.8 = s32[]{:T(128)} parameter(3), metadata={op_name="top_k"} - %convert.385 = f32[]{:T(128)S(6)} convert(%top_k.0), metadata={op_name="convert.16"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} - %bitcast-convert.35 = s32[]{:T(128)S(6)} bitcast-convert(%convert.385), metadata={op_name="bitcast-convert.6"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} - %compare.128 = pred[]{:T(512)S(6)} compare(%bitcast-convert.35, %constant.1387), direction=LT, metadata={op_name="compare.35"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} - %xor.36 = s32[]{:T(128)S(6)} xor(%constant.1388, %bitcast-convert.35), metadata={op_name="xor.6"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %convert.384 = f32[]{:T(128)S(6)} convert(%top_k.0), metadata={op_name="convert.16"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %bitcast-convert.35 = s32[]{:T(128)S(6)} bitcast-convert(%convert.384), metadata={op_name="bitcast-convert.6"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %compare.128 = pred[]{:T(512)S(6)} compare(%bitcast-convert.35, %constant.1521), direction=LT, metadata={op_name="compare.35"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %xor.36 = s32[]{:T(128)S(6)} xor(%constant.1522, %bitcast-convert.35), metadata={op_name="xor.6"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} %select.118 = s32[]{:T(128)S(6)} select(%compare.128, %xor.36, %bitcast-convert.35), metadata={op_name="select.14"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[{"indices":["1","3"]}]}} - %convert.386 = f32[]{:T(128)S(6)} convert(%top_k.6), metadata={op_name="convert.17"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} - %bitcast-convert.36 = s32[]{:T(128)S(6)} bitcast-convert(%convert.386), metadata={op_name="bitcast-convert.7"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} - %compare.129 = pred[]{:T(512)S(6)} compare(%bitcast-convert.36, %constant.1387), direction=LT, metadata={op_name="compare.36"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} - %xor.37 = s32[]{:T(128)S(6)} xor(%constant.1388, %bitcast-convert.36), metadata={op_name="xor.7"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %convert.385 = f32[]{:T(128)S(6)} convert(%top_k.6), metadata={op_name="convert.17"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %bitcast-convert.36 = s32[]{:T(128)S(6)} bitcast-convert(%convert.385), metadata={op_name="bitcast-convert.7"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %compare.129 = pred[]{:T(512)S(6)} compare(%bitcast-convert.36, %constant.1521), direction=LT, metadata={op_name="compare.36"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %xor.37 = s32[]{:T(128)S(6)} xor(%constant.1522, %bitcast-convert.36), metadata={op_name="xor.7"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} %select.119 = s32[]{:T(128)S(6)} select(%compare.129, %xor.37, %bitcast-convert.36), metadata={op_name="select.15"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[{"indices":["1","3"]}]}} %compare.130 = pred[]{:T(512)S(6)} compare(%select.118, %select.119), direction=GT, metadata={op_name="compare.1"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} %compare.131 = pred[]{:T(512)S(6)} compare(%select.119, %select.118), direction=GT, metadata={op_name="compare.108"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} - %compare.132 = pred[]{:T(512)S(6)} compare(%compare.130, %compare.131), direction=EQ, metadata={op_name="compare.109"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} - %compare.133 = pred[]{:T(512)S(6)} compare(%top_k.7, %top_k.8), direction=LT, metadata={op_name="compare.110"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} - ROOT %select.120 = pred[]{:T(512)} select(%compare.132, %compare.133, %compare.130), metadata={op_name="select.108"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_15.21.clone.1 (reduce-window.326: s32[], reduce-window.327: s32[]) -> s32[] { - %reduce-window.326 = s32[]{:T(128)} parameter(0), metadata={op_name="reduce-window.20"} - %reduce-window.327 = s32[]{:T(128)} parameter(1), metadata={op_name="reduce-window.20"} - ROOT %reduce_window_sum.282 = s32[]{:T(128)} add(%reduce-window.326, %reduce-window.327), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_16.22.clone.1 (reduce-window.330: s32[], reduce-window.331: s32[]) -> s32[] { - %reduce-window.330 = s32[]{:T(128)} parameter(0), metadata={op_name="reduce-window.56"} - %reduce-window.331 = s32[]{:T(128)} parameter(1), metadata={op_name="reduce-window.56"} - ROOT %reduce_window_sum.284 = s32[]{:T(128)} add(%reduce-window.330, %reduce-window.331), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_18.24.clone.1 (reduce-window.334: s32[], reduce-window.335: s32[]) -> s32[] { - %reduce-window.334 = s32[]{:T(128)} parameter(0), metadata={op_name="reduce-window.22"} - %reduce-window.335 = s32[]{:T(128)} parameter(1), metadata={op_name="reduce-window.22"} - ROOT %reduce_window_sum.286 = s32[]{:T(128)} add(%reduce-window.334, %reduce-window.335), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_19.25.clone.1 (reduce-window.338: s32[], reduce-window.339: s32[]) -> s32[] { - %reduce-window.338 = s32[]{:T(128)} parameter(0), metadata={op_name="reduce-window.57"} - %reduce-window.339 = s32[]{:T(128)} parameter(1), metadata={op_name="reduce-window.57"} - ROOT %reduce_window_sum.288 = s32[]{:T(128)} add(%reduce-window.338, %reduce-window.339), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_21.27.clone.1 (reduce-window.342: s32[], reduce-window.343: s32[]) -> s32[] { - %reduce-window.342 = s32[]{:T(128)} parameter(0), metadata={op_name="reduce-window.24"} - %reduce-window.343 = s32[]{:T(128)} parameter(1), metadata={op_name="reduce-window.24"} - ROOT %reduce_window_sum.290 = s32[]{:T(128)} add(%reduce-window.342, %reduce-window.343), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_22.28.clone.1 (reduce-window.346: s32[], reduce-window.347: s32[]) -> s32[] { - %reduce-window.346 = s32[]{:T(128)} parameter(0), metadata={op_name="reduce-window.58"} - %reduce-window.347 = s32[]{:T(128)} parameter(1), metadata={op_name="reduce-window.58"} - ROOT %reduce_window_sum.292 = s32[]{:T(128)} add(%reduce-window.346, %reduce-window.347), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.4.clone (param_0.68: s32[256], param_1.114: s32[1024]) -> s32[263] { - %param_0.68 = s32[256]{0:T(256)S(1)} parameter(0) - %param_1.114 = s32[1024]{0:T(1024)S(1)} parameter(1) - %custom-call.19 = s32[1024]{0:T(1024)} custom-call(%param_1.114), custom_call_target="AssumeGatherIndicesInBound", operand_layout_constraints={s32[1024]{0:T(1024)}}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/jit(_take)/gather" stack_frame_id=0} - %slice.898 = s32[263]{0:T(512)} slice(%custom-call.19), slice={[0:263]}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/jit(_take)/gather" stack_frame_id=0} - %reshape.3622 = s32[263]{0:T(512)} reshape(%slice.898), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/jit(_take)/broadcast_in_dim" stack_frame_id=0} - %transpose.1037 = s32[263]{0:T(512)} transpose(%reshape.3622), dimensions={0}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/jit(_take)/broadcast_in_dim" stack_frame_id=0} - %gather.209 = s32[263]{0:T(512)} gather(%param_0.68, %transpose.1037), offset_dims={}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/jit(_take)/gather" stack_frame_id=0} - %transpose.1036 = s32[263]{0:T(512)} transpose(%gather.209), dimensions={0}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/jit(_take)/gather" stack_frame_id=0} - ROOT %reshape.3621 = s32[263]{0:T(512)S(1)} reshape(%transpose.1036), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/jit(_take)/gather" stack_frame_id=0} -} - -%region_27.34.clone.1 (reduce-window.350: s32[], reduce-window.351: s32[]) -> s32[] { - %reduce-window.350 = s32[]{:T(128)} parameter(0), metadata={op_name="reduce-window.26"} - %reduce-window.351 = s32[]{:T(128)} parameter(1), metadata={op_name="reduce-window.26"} - ROOT %reduce_window_sum.294 = s32[]{:T(128)} add(%reduce-window.350, %reduce-window.351), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_29.36.clone.1 (reduce-window.354: s32[], reduce-window.355: s32[]) -> s32[] { - %reduce-window.354 = s32[]{:T(128)} parameter(0), metadata={op_name="reduce-window.27"} - %reduce-window.355 = s32[]{:T(128)} parameter(1), metadata={op_name="reduce-window.27"} - ROOT %reduce_window_sum.296 = s32[]{:T(128)} add(%reduce-window.354, %reduce-window.355), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_30.37.clone.1 (reduce-window.358: s32[], reduce-window.359: s32[]) -> s32[] { - %reduce-window.358 = s32[]{:T(128)} parameter(0), metadata={op_name="reduce-window.59"} - %reduce-window.359 = s32[]{:T(128)} parameter(1), metadata={op_name="reduce-window.59"} - ROOT %reduce_window_sum.298 = s32[]{:T(128)} add(%reduce-window.358, %reduce-window.359), metadata={op_name="reduce_window_sum"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_13.19 (sort.44: s32[], sort.45: s32[], sort.46: s32[], sort.47: s32[], sort.48: s32[], sort.49: s32[]) -> pred[] { - %sort.46 = s32[]{:T(128)} parameter(2), metadata={op_name="jit(argsort)/sort"} - %sort.47 = s32[]{:T(128)} parameter(3), metadata={op_name="jit(argsort)/sort"} - %sort.44 = s32[]{:T(128)} parameter(0), metadata={op_name="jit(argsort)/sort"} - %sort.45 = s32[]{:T(128)} parameter(1), metadata={op_name="jit(argsort)/sort"} - %sort.48 = s32[]{:T(128)} parameter(4), metadata={op_name="jit(argsort)/sort"} - %sort.49 = s32[]{:T(128)} parameter(5), metadata={op_name="jit(argsort)/sort"} - %lt_to.27 = pred[]{:T(512)S(6)} compare(%sort.44, %sort.45), direction=LT, metadata={op_name="lt_to"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} - %lt_to.28 = pred[]{:T(512)S(6)} compare(%sort.45, %sort.44), direction=LT, metadata={op_name="lt_to"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} - %compare.134 = pred[]{:T(512)S(6)} compare(%lt_to.27, %lt_to.28), direction=EQ, metadata={op_name="compare.111"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} - %compare.135 = pred[]{:T(512)S(6)} compare(%sort.48, %sort.49), direction=LT, metadata={op_name="compare.112"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} - ROOT %select.121 = pred[]{:T(512)} select(%compare.134, %compare.135, %lt_to.27), metadata={op_name="select.109"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.2.clone (param_0.71: bf16[4096,512], param_1.116: s32[4096]) -> bf16[4096,512] { - %param_0.71 = bf16[4096,512]{1,0:T(8,128)(2,1)S(1)} parameter(0) - %param_1.116 = s32[4096]{0:T(1024)S(1)} parameter(1) - %custom-call.21 = s32[4096]{0:T(1024)} custom-call(%param_1.116), custom_call_target="AssumeGatherIndicesInBound", operand_layout_constraints={s32[4096]{0:T(1024)}}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} - %slice.900 = s32[4096]{0:T(1024)} slice(%custom-call.21), slice={[0:4096]}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} - %reshape.3645 = s32[4096]{0:T(1024)} reshape(%slice.900), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/sort_activations/broadcast_in_dim" stack_frame_id=0} - %transpose.1043 = s32[4096]{0:T(1024)} transpose(%reshape.3645), dimensions={0}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/sort_activations/broadcast_in_dim" stack_frame_id=0} - %gather.210 = bf16[4096,512]{1,0:T(8,128)(2,1)} gather(%param_0.71, %transpose.1043), offset_dims={1}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1,512}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} - %transpose.1042 = bf16[4096,512]{1,0:T(8,128)(2,1)} transpose(%gather.210), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} - ROOT %reshape.3644 = bf16[4096,512]{1,0:T(8,128)(2,1)S(1)} reshape(%transpose.1042), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} -} - -%region_31.39 (sort.50: s32[], sort.51: s32[], sort.52: s32[], sort.53: s32[], sort.54: s32[], sort.55: s32[]) -> pred[] { - %sort.52 = s32[]{:T(128)} parameter(2), metadata={op_name="jit(argsort)/sort"} - %sort.53 = s32[]{:T(128)} parameter(3), metadata={op_name="jit(argsort)/sort"} - %sort.50 = s32[]{:T(128)} parameter(0), metadata={op_name="jit(argsort)/sort"} - %sort.51 = s32[]{:T(128)} parameter(1), metadata={op_name="jit(argsort)/sort"} - %sort.54 = s32[]{:T(128)} parameter(4), metadata={op_name="jit(argsort)/sort"} - %sort.55 = s32[]{:T(128)} parameter(5), metadata={op_name="jit(argsort)/sort"} - %lt_to.30 = pred[]{:T(512)S(6)} compare(%sort.50, %sort.51), direction=LT, metadata={op_name="lt_to"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} - %lt_to.31 = pred[]{:T(512)S(6)} compare(%sort.51, %sort.50), direction=LT, metadata={op_name="lt_to"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} - %compare.142 = pred[]{:T(512)S(6)} compare(%lt_to.30, %lt_to.31), direction=EQ, metadata={op_name="compare.113"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} - %compare.143 = pred[]{:T(512)S(6)} compare(%sort.54, %sort.55), direction=LT, metadata={op_name="compare.114"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} - ROOT %select.126 = pred[]{:T(512)} select(%compare.142, %compare.143, %lt_to.30), metadata={op_name="select.110"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.3.clone (param_0.72: bf16[4096,512], param_1.118: s32[4096]) -> bf16[4096,512] { - %param_0.72 = bf16[4096,512]{1,0:T(8,128)(2,1)S(1)} parameter(0) - %param_1.118 = s32[4096]{0:T(1024)S(1)} parameter(1) - %custom-call.23 = s32[4096]{0:T(1024)} custom-call(%param_1.118), custom_call_target="AssumeGatherIndicesInBound", operand_layout_constraints={s32[4096]{0:T(1024)}}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} - %slice.902 = s32[4096]{0:T(1024)} slice(%custom-call.23), slice={[0:4096]}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} - %reshape.3647 = s32[4096]{0:T(1024)} reshape(%slice.902), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/sort_activations/broadcast_in_dim" stack_frame_id=0} - %transpose.1045 = s32[4096]{0:T(1024)} transpose(%reshape.3647), dimensions={0}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/sort_activations/broadcast_in_dim" stack_frame_id=0} - %gather.211 = bf16[4096,512]{1,0:T(8,128)(2,1)} gather(%param_0.72, %transpose.1045), offset_dims={1}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=1, slice_sizes={1,512}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} - %transpose.1044 = bf16[4096,512]{1,0:T(8,128)(2,1)} transpose(%gather.211), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} - ROOT %reshape.3646 = bf16[4096,512]{1,0:T(8,128)(2,1)S(1)} reshape(%transpose.1044), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/sort_activations/gather" stack_frame_id=0} -} - -%compare (name: s32[], name.1: s32[], name.2: bf16[], name.3: bf16[]) -> pred[] { - %name.2 = bf16[] parameter(2) - %name.3 = bf16[] parameter(3) - %name = s32[] parameter(0) - %name.1 = s32[] parameter(1) - ROOT %compare.377 = pred[] compare(%name, %name.1), direction=LT, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%compare.1 (name.4: s32[], name.5: s32[], name.6: f32[], name.7: f32[]) -> pred[] { - %name.6 = f32[] parameter(2) - %name.7 = f32[] parameter(3) - %name.4 = s32[] parameter(0) - %name.5 = s32[] parameter(1) - ROOT %compare.378 = pred[] compare(%name.4, %name.5), direction=LT, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%compare.2 (name.8: s32[], name.9: s32[], name.10: f32[], name.11: f32[]) -> pred[] { - %name.10 = f32[] parameter(2) - %name.11 = f32[] parameter(3) - %name.8 = s32[] parameter(0) - %name.9 = s32[] parameter(1) - ROOT %compare.379 = pred[] compare(%name.8, %name.9), direction=LT, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%compare.3 (name.12: s32[], name.13: s32[], name.14: f32[], name.15: f32[]) -> pred[] { - %name.14 = f32[] parameter(2) - %name.15 = f32[] parameter(3) - %name.12 = s32[] parameter(0) - %name.13 = s32[] parameter(1) - ROOT %compare.380 = pred[] compare(%name.12, %name.13), direction=LT, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%compare.4 (name.16: s32[], name.17: s32[], name.18: f32[], name.19: f32[]) -> pred[] { - %name.18 = f32[] parameter(2) - %name.19 = f32[] parameter(3) - %name.16 = s32[] parameter(0) - %name.17 = s32[] parameter(1) - ROOT %compare.381 = pred[] compare(%name.16, %name.17), direction=LT, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%called_computation.13 (param_0.4519: s32[256]) -> s32[256] { - %param_0.4519 = s32[256]{0:T(256)} parameter(0) - ROOT %copy.2073 = s32[256]{0:T(256)} copy(%param_0.4519), backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["16"],"input_window_bounds":[],"estimated_cycles":"1134","iteration_bounds":[],"scratchpad_allocation_size":"256","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"16","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} -}, execution_thread="sparsecore" - -%async_computation.13 (param_0.4520: s32[256]) -> s32[256] { - %param_0.4520 = s32[256]{0:T(256)} parameter(0) - ROOT %copy.2074.cloned.1 = s32[256]{0:T(256)} call(%param_0.4520), to_apply=%called_computation.13 -}, execution_thread="sparsecore" - -%region_49.59 (scatter-add.14: s32[], scatter-add.15: s32[]) -> s32[] { - %scatter-add.14 = s32[]{:T(128)S(7)} parameter(0), metadata={op_name="scatter-add"} - %scatter-add.15 = s32[]{:T(128)S(7)} parameter(1), metadata={op_name="scatter-add"} - ROOT %add.1387 = s32[]{:T(128)S(7)} add(%scatter-add.14, %scatter-add.15), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} -}, execution_thread="sparsecore" - -%fused_computation.22.clone.clone (param_0.4521: s32[256], param_1.5326: s32[4096], param_2.4491: s32[4096]) -> s32[256] { - %param_0.4521 = s32[256]{0:T(256)} parameter(0) - %param_1.5326 = s32[4096]{0:T(1024)} parameter(1) - %reshape.3911 = s32[4096]{0:T(1024)} reshape(%param_1.5326), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(clip)/max" stack_frame_id=0} - %transpose.1100 = s32[4096]{0:T(1024)} transpose(%reshape.3911), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(clip)/max" stack_frame_id=0} - %param_2.4491 = s32[4096]{0:T(1024)} parameter(2) - %reshape.3912 = s32[4096]{0:T(1024)} reshape(%param_2.4491), metadata={op_name="jit(train_step)/moe_layers/shard_map/broadcast_in_dim" stack_frame_id=0} - %transpose.1101 = s32[4096]{0:T(1024)} transpose(%reshape.3912), dimensions={0}, metadata={op_name="jit(train_step)/moe_layers/shard_map/broadcast_in_dim" stack_frame_id=0} - ROOT %scatter-add.231 = s32[256]{0:T(256)} scatter(%param_0.4521, %transpose.1100, %transpose.1101), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, to_apply=%region_49.59, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/scatter-add" stack_frame_id=0} -}, execution_thread="sparsecore" - -%called_computation.14 (param_0.4522: s32[256], param_1.5327: s32[4096], param_2.4492: s32[4096]) -> s32[256] { - %param_0.4522 = s32[256]{0:T(256)} parameter(0) - %param_1.5327 = s32[4096]{0:T(1024)} parameter(1) - %param_2.4492 = s32[4096]{0:T(1024)} parameter(2) - ROOT %scatter_offload_custom_fusion.39 = s32[256]{0:T(256)} fusion(%param_0.4522, %param_1.5327, %param_2.4492), kind=kCustom, calls=%fused_computation.22.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["256"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"4160","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} -}, execution_thread="sparsecore" - -%async_computation.14 (param_0.4523: s32[256], param_1.5328: s32[4096], param_2.4493: s32[4096]) -> s32[256] { - %param_0.4523 = s32[256]{0:T(256)} parameter(0) - %param_1.5328 = s32[4096]{0:T(1024)} parameter(1) - %param_2.4493 = s32[4096]{0:T(1024)} parameter(2) - ROOT %scatter_offload_custom_fusion.40.cloned.1 = s32[256]{0:T(256)} call(%param_0.4523, %param_1.5328, %param_2.4493), to_apply=%called_computation.14, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/scatter-add" stack_frame_id=0} -}, execution_thread="sparsecore" - -%called_computation (param_0.84: s32[256], param_1.136: s32[4096], param_2.80: s32[4096], param_3.3085: token[]) -> s32[256] { - %param_3.3085 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} - %param_0.84 = s32[256]{0:T(256)} parameter(0), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} - %param_1.136 = s32[4096]{0:T(1024)} parameter(1), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} - %param_2.80 = s32[4096]{0:T(1024)} parameter(2), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} - %copy.2074.cloned.1.call-start = ((s32[256]{0:T(256)}), s32[256]{0:T(256)}, u32[]{:S(8)}) async-start(%param_0.84), async_execution_thread="sparsecore", calls=%async_computation.13 - %copy.2074.cloned.1.call-done = s32[256]{0:T(256)} async-done(%copy.2074.cloned.1.call-start) - %scatter_offload_custom_fusion.40.cloned.1.call-start = ((s32[256]{0:T(256)}, s32[4096]{0:T(1024)}, s32[4096]{0:T(1024)}), s32[256]{0:T(256)}, u32[]{:S(8)}) async-start(%copy.2074.cloned.1.call-done, %param_1.136, %param_2.80), async_execution_thread="sparsecore", calls=%async_computation.14, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/scatter-add" stack_frame_id=0} - ROOT %scatter_offload_custom_fusion.40.cloned.1.call-done = s32[256]{0:T(256)} async-done(%scatter_offload_custom_fusion.40.cloned.1.call-start), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/scatter-add" stack_frame_id=0} -}, execution_thread="sparsecore" - -%async_computation (param_0.85: s32[256], param_1.137: s32[4096], param_2.81: s32[4096], param_3.3084: token[]) -> s32[256] { - %param_3.3084 = token[] parameter(3) - %param_0.85 = s32[256]{0:T(256)} parameter(0) - %param_1.137 = s32[4096]{0:T(1024)} parameter(1) - %param_2.81 = s32[4096]{0:T(1024)} parameter(2) - ROOT %scatter_offload_custom_fusion.2.cloned.1 = s32[256]{0:T(256)} call(%param_0.85, %param_1.137, %param_2.81, %param_3.3084), to_apply=%called_computation, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/scatter-add" stack_frame_id=0} -}, execution_thread="sparsecore" - -%called_computation.15 (param_0.4524: f32[9]) -> f32[9] { - %param_0.4524 = f32[9]{0:T(128)} parameter(0) - ROOT %copy.2075 = f32[9]{0:T(128)} copy(%param_0.4524), backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["8"],"input_window_bounds":[],"estimated_cycles":"1131","iteration_bounds":[],"scratchpad_allocation_size":"128","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"16","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} -}, execution_thread="sparsecore" - -%async_computation.15 (param_0.4525: f32[9]) -> f32[9] { - %param_0.4525 = f32[9]{0:T(128)} parameter(0) - ROOT %copy.2076.cloned.1 = f32[9]{0:T(128)} call(%param_0.4525), to_apply=%called_computation.15 -}, execution_thread="sparsecore" - -%region_61.72 (scatter-add.24: f32[], scatter-add.25: f32[]) -> f32[] { - %scatter-add.24 = f32[]{:T(128)S(7)} parameter(0), metadata={op_name="scatter-add"} - %scatter-add.25 = f32[]{:T(128)S(7)} parameter(1), metadata={op_name="scatter-add"} - ROOT %add.1393 = f32[]{:T(128)S(7)} add(%scatter-add.24, %scatter-add.25), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} -}, execution_thread="sparsecore" - -%fused_computation.24.clone.clone (param_0.4526: f32[9], param_1.5329: s32[256], param_2.4494: f32[256]) -> f32[9] { - %param_0.4526 = f32[9]{0:T(128)} parameter(0) - %param_1.5329 = s32[256]{0:T(256)} parameter(1) - %reshape.3913 = s32[256]{0:T(256)} reshape(%param_1.5329), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} - %transpose.1102 = s32[256]{0:T(256)} transpose(%reshape.3913), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} - %param_2.4494 = f32[256]{0:T(256)} parameter(2) - %reshape.3914 = f32[256]{0:T(256)} reshape(%param_2.4494), metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} - %transpose.1103 = f32[256]{0:T(256)} transpose(%reshape.3914), dimensions={0}, metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} - ROOT %scatter-add.232 = f32[9]{0:T(128)} scatter(%param_0.4526, %transpose.1102, %transpose.1103), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, indices_are_sorted=true, to_apply=%region_61.72, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} -}, execution_thread="sparsecore" - -%called_computation.16 (param_0.4527: f32[9], param_1.5330: s32[256], param_2.4495: f32[256]) -> f32[9] { - %param_0.4527 = f32[9]{0:T(128)} parameter(0) - %param_1.5330 = s32[256]{0:T(256)} parameter(1) - %param_2.4495 = f32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.41 = f32[9]{0:T(128)} fusion(%param_0.4527, %param_1.5330, %param_2.4495), kind=kCustom, calls=%fused_computation.24.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["16"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"1312","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} -}, execution_thread="sparsecore" - -%async_computation.16 (param_0.4528: f32[9], param_1.5331: s32[256], param_2.4496: f32[256]) -> f32[9] { - %param_0.4528 = f32[9]{0:T(128)} parameter(0) - %param_1.5331 = s32[256]{0:T(256)} parameter(1) - %param_2.4496 = f32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.42.cloned.1 = f32[9]{0:T(128)} call(%param_0.4528, %param_1.5331, %param_2.4496), to_apply=%called_computation.16, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} -}, execution_thread="sparsecore" - -%called_computation.1 (param_0.87: f32[9], param_1.139: s32[256], param_2.83: f32[256], param_3.3099: token[]) -> f32[9] { - %param_3.3099 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} - %param_0.87 = f32[9]{0:T(128)} parameter(0), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} - %param_1.139 = s32[256]{0:T(256)} parameter(1), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} - %param_2.83 = f32[256]{0:T(256)} parameter(2), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} - %copy.2076.cloned.1.call-start = ((f32[9]{0:T(128)}), f32[9]{0:T(128)}, u32[]{:S(8)}) async-start(%param_0.87), async_execution_thread="sparsecore", calls=%async_computation.15 - %copy.2076.cloned.1.call-done = f32[9]{0:T(128)} async-done(%copy.2076.cloned.1.call-start) - %scatter_offload_custom_fusion.42.cloned.1.call-start = ((f32[9]{0:T(128)}, s32[256]{0:T(256)}, f32[256]{0:T(256)}), f32[9]{0:T(128)}, u32[]{:S(8)}) async-start(%copy.2076.cloned.1.call-done, %param_1.139, %param_2.83), async_execution_thread="sparsecore", calls=%async_computation.16, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} - ROOT %scatter_offload_custom_fusion.42.cloned.1.call-done = f32[9]{0:T(128)} async-done(%scatter_offload_custom_fusion.42.cloned.1.call-start), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} -}, execution_thread="sparsecore" - -%async_computation.1 (param_0.88: f32[9], param_1.140: s32[256], param_2.84: f32[256], param_3.3098: token[]) -> f32[9] { - %param_3.3098 = token[] parameter(3) - %param_0.88 = f32[9]{0:T(128)} parameter(0) - %param_1.140 = s32[256]{0:T(256)} parameter(1) - %param_2.84 = f32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.5.cloned.1 = f32[9]{0:T(128)} call(%param_0.88, %param_1.140, %param_2.84, %param_3.3098), to_apply=%called_computation.1, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} -}, execution_thread="sparsecore" - -%called_computation.17 (param_0.4529: s32[263]) -> s32[263] { - %param_0.4529 = s32[263]{0:T(512)} parameter(0) - ROOT %copy.2077 = s32[263]{0:T(512)} copy(%param_0.4529), backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["32"],"input_window_bounds":[],"estimated_cycles":"1141","iteration_bounds":[],"scratchpad_allocation_size":"512","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"16","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} -}, execution_thread="sparsecore" - -%async_computation.17 (param_0.4530: s32[263]) -> s32[263] { - %param_0.4530 = s32[263]{0:T(512)} parameter(0) - ROOT %copy.2078.cloned.1 = s32[263]{0:T(512)} call(%param_0.4530), to_apply=%called_computation.17 -}, execution_thread="sparsecore" - -%region_63.74 (scatter-add.28: s32[], scatter-add.29: s32[]) -> s32[] { - %scatter-add.28 = s32[]{:T(128)S(7)} parameter(0), metadata={op_name="scatter-add"} - %scatter-add.29 = s32[]{:T(128)S(7)} parameter(1), metadata={op_name="scatter-add"} - ROOT %add.1394 = s32[]{:T(128)S(7)} add(%scatter-add.28, %scatter-add.29), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} -}, execution_thread="sparsecore" - -%fused_computation.25.clone.clone (param_0.4531: s32[263], param_1.5332: s32[8], param_2.4497: s32[8]) -> s32[263] { - %param_0.4531 = s32[263]{0:T(512)} parameter(0) - %param_1.5332 = s32[8]{0:T(128)} parameter(1) - %reshape.3915 = s32[8]{0:T(128)} reshape(%param_1.5332), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/select_n" stack_frame_id=0} - %transpose.1104 = s32[8]{0:T(128)} transpose(%reshape.3915), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/select_n" stack_frame_id=0} - %param_2.4497 = s32[8]{0:T(128)} parameter(2) - %reshape.3916 = s32[8]{0:T(128)} reshape(%param_2.4497), metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast.80" stack_frame_id=0} - %transpose.1105 = s32[8]{0:T(128)} transpose(%reshape.3916), dimensions={0}, metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast.80" stack_frame_id=0} - ROOT %scatter-add.233 = s32[263]{0:T(512)} scatter(%param_0.4531, %transpose.1104, %transpose.1105), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, to_apply=%region_63.74, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} -}, execution_thread="sparsecore" - -%called_computation.18 (param_0.4532: s32[263], param_1.5333: s32[8], param_2.4498: s32[8]) -> s32[263] { - %param_0.4532 = s32[263]{0:T(512)} parameter(0) - %param_1.5333 = s32[8]{0:T(128)} parameter(1) - %param_2.4498 = s32[8]{0:T(128)} parameter(2) - ROOT %scatter_offload_custom_fusion.43 = s32[263]{0:T(512)} fusion(%param_0.4532, %param_1.5333, %param_2.4498), kind=kCustom, calls=%fused_computation.25.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["8"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"256","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} -}, execution_thread="sparsecore" - -%async_computation.18 (param_0.4533: s32[263], param_1.5334: s32[8], param_2.4499: s32[8]) -> s32[263] { - %param_0.4533 = s32[263]{0:T(512)} parameter(0) - %param_1.5334 = s32[8]{0:T(128)} parameter(1) - %param_2.4499 = s32[8]{0:T(128)} parameter(2) - ROOT %scatter_offload_custom_fusion.44.cloned.1 = s32[263]{0:T(512)} call(%param_0.4533, %param_1.5334, %param_2.4499), to_apply=%called_computation.18, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} -}, execution_thread="sparsecore" - -%called_computation.2 (param_0.90: s32[263], param_1.142: s32[8], param_2.86: s32[8], param_3.3105: token[]) -> s32[263] { - %param_3.3105 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} - %param_0.90 = s32[263]{0:T(512)} parameter(0), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} - %param_1.142 = s32[8]{0:T(128)} parameter(1), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} - %param_2.86 = s32[8]{0:T(128)} parameter(2), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} - %copy.2078.cloned.1.call-start = ((s32[263]{0:T(512)}), s32[263]{0:T(512)}, u32[]{:S(8)}) async-start(%param_0.90), async_execution_thread="sparsecore", calls=%async_computation.17 - %copy.2078.cloned.1.call-done = s32[263]{0:T(512)} async-done(%copy.2078.cloned.1.call-start) - %scatter_offload_custom_fusion.44.cloned.1.call-start = ((s32[263]{0:T(512)}, s32[8]{0:T(128)}, s32[8]{0:T(128)}), s32[263]{0:T(512)}, u32[]{:S(8)}) async-start(%copy.2078.cloned.1.call-done, %param_1.142, %param_2.86), async_execution_thread="sparsecore", calls=%async_computation.18, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} - ROOT %scatter_offload_custom_fusion.44.cloned.1.call-done = s32[263]{0:T(512)} async-done(%scatter_offload_custom_fusion.44.cloned.1.call-start), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} -}, execution_thread="sparsecore" - -%async_computation.2 (param_0.91: s32[263], param_1.143: s32[8], param_2.87: s32[8], param_3.3104: token[]) -> s32[263] { - %param_3.3104 = token[] parameter(3) - %param_0.91 = s32[263]{0:T(512)} parameter(0) - %param_1.143 = s32[8]{0:T(128)} parameter(1) - %param_2.87 = s32[8]{0:T(128)} parameter(2) - ROOT %scatter_offload_custom_fusion.8.cloned.1 = s32[263]{0:T(512)} call(%param_0.91, %param_1.143, %param_2.87, %param_3.3104), to_apply=%called_computation.2, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/rematted_computation/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} -}, execution_thread="sparsecore" - -%called_computation.19 (param_0.4534: s32[263]) -> s32[263] { - %param_0.4534 = s32[263]{0:T(512)} parameter(0) - ROOT %copy.2079 = s32[263]{0:T(512)} copy(%param_0.4534), backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["32"],"input_window_bounds":[],"estimated_cycles":"1141","iteration_bounds":[],"scratchpad_allocation_size":"512","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"16","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} -}, execution_thread="sparsecore" - -%async_computation.19 (param_0.4535: s32[263]) -> s32[263] { - %param_0.4535 = s32[263]{0:T(512)} parameter(0) - ROOT %copy.2080.cloned.1 = s32[263]{0:T(512)} call(%param_0.4535), to_apply=%called_computation.19 -}, execution_thread="sparsecore" - -%region_73.86.clone (scatter-add.163: s32[], scatter-add.164: s32[]) -> s32[] { - %scatter-add.163 = s32[]{:T(128)S(7)} parameter(0), metadata={op_name="scatter-add"} - %scatter-add.164 = s32[]{:T(128)S(7)} parameter(1), metadata={op_name="scatter-add"} - ROOT %add.2527 = s32[]{:T(128)S(7)} add(%scatter-add.163, %scatter-add.164), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} -}, execution_thread="sparsecore" - -%fused_computation.26.clone.clone (param_0.4536: s32[263], param_1.5335: s32[256], param_2.4500: s32[256]) -> s32[263] { - %param_0.4536 = s32[263]{0:T(512)} parameter(0) - %param_1.5335 = s32[256]{0:T(256)} parameter(1) - %reshape.3917 = s32[256]{0:T(256)} reshape(%param_1.5335), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/select_n" stack_frame_id=0} - %transpose.1106 = s32[256]{0:T(256)} transpose(%reshape.3917), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/select_n" stack_frame_id=0} - %param_2.4500 = s32[256]{0:T(256)} parameter(2) - %reshape.3918 = s32[256]{0:T(256)} reshape(%param_2.4500), metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} - %transpose.1107 = s32[256]{0:T(256)} transpose(%reshape.3918), dimensions={0}, metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} - ROOT %scatter-add.234 = s32[263]{0:T(512)} scatter(%param_0.4536, %transpose.1106, %transpose.1107), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, to_apply=%region_73.86.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} -}, execution_thread="sparsecore" - -%called_computation.20 (param_0.4537: s32[263], param_1.5336: s32[256], param_2.4501: s32[256]) -> s32[263] { - %param_0.4537 = s32[263]{0:T(512)} parameter(0) - %param_1.5336 = s32[256]{0:T(256)} parameter(1) - %param_2.4501 = s32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.45 = s32[263]{0:T(512)} fusion(%param_0.4537, %param_1.5336, %param_2.4501), kind=kCustom, calls=%fused_computation.26.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["16"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"384","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} -}, execution_thread="sparsecore" - -%async_computation.20 (param_0.4538: s32[263], param_1.5337: s32[256], param_2.4502: s32[256]) -> s32[263] { - %param_0.4538 = s32[263]{0:T(512)} parameter(0) - %param_1.5337 = s32[256]{0:T(256)} parameter(1) - %param_2.4502 = s32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.46.cloned.1 = s32[263]{0:T(512)} call(%param_0.4538, %param_1.5337, %param_2.4502), to_apply=%called_computation.20, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} -}, execution_thread="sparsecore" - -%called_computation.3 (param_0.93: s32[263], param_1.145: s32[256], param_2.89: s32[256], param_3.3091: token[]) -> s32[263] { - %param_3.3091 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} - %param_0.93 = s32[263]{0:T(512)} parameter(0), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} - %param_1.145 = s32[256]{0:T(256)} parameter(1), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} - %param_2.89 = s32[256]{0:T(256)} parameter(2), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} - %copy.2080.cloned.1.call-start = ((s32[263]{0:T(512)}), s32[263]{0:T(512)}, u32[]{:S(8)}) async-start(%param_0.93), async_execution_thread="sparsecore", calls=%async_computation.19 - %copy.2080.cloned.1.call-done = s32[263]{0:T(512)} async-done(%copy.2080.cloned.1.call-start) - %scatter_offload_custom_fusion.46.cloned.1.call-start = ((s32[263]{0:T(512)}, s32[256]{0:T(256)}, s32[256]{0:T(256)}), s32[263]{0:T(512)}, u32[]{:S(8)}) async-start(%copy.2080.cloned.1.call-done, %param_1.145, %param_2.89), async_execution_thread="sparsecore", calls=%async_computation.20, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} - ROOT %scatter_offload_custom_fusion.46.cloned.1.call-done = s32[263]{0:T(512)} async-done(%scatter_offload_custom_fusion.46.cloned.1.call-start), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} -}, execution_thread="sparsecore" - -%async_computation.3 (param_0.94: s32[263], param_1.146: s32[256], param_2.90: s32[256], param_3.3090: token[]) -> s32[263] { - %param_3.3090 = token[] parameter(3) - %param_0.94 = s32[263]{0:T(512)} parameter(0) - %param_1.146 = s32[256]{0:T(256)} parameter(1) - %param_2.90 = s32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.11.cloned.1 = s32[263]{0:T(512)} call(%param_0.94, %param_1.146, %param_2.90, %param_3.3090), to_apply=%called_computation.3, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} -}, execution_thread="sparsecore" - -%called_computation.21 (param_0.4539: f32[9]) -> f32[9] { - %param_0.4539 = f32[9]{0:T(128)} parameter(0) - ROOT %copy.2081 = f32[9]{0:T(128)} copy(%param_0.4539), backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["8"],"input_window_bounds":[],"estimated_cycles":"1131","iteration_bounds":[],"scratchpad_allocation_size":"128","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"16","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} -}, execution_thread="sparsecore" - -%async_computation.21 (param_0.4540: f32[9]) -> f32[9] { - %param_0.4540 = f32[9]{0:T(128)} parameter(0) - ROOT %copy.2082.cloned.1 = f32[9]{0:T(128)} call(%param_0.4540), to_apply=%called_computation.21 -}, execution_thread="sparsecore" - -%region_79.95.clone (scatter-add.167: f32[], scatter-add.168: f32[]) -> f32[] { - %scatter-add.167 = f32[]{:T(128)S(7)} parameter(0), metadata={op_name="scatter-add"} - %scatter-add.168 = f32[]{:T(128)S(7)} parameter(1), metadata={op_name="scatter-add"} - ROOT %add.2529 = f32[]{:T(128)S(7)} add(%scatter-add.167, %scatter-add.168), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} -}, execution_thread="sparsecore" - -%fused_computation.27.clone.clone (param_0.4541: f32[9], param_1.5338: s32[256], param_2.4503: f32[256]) -> f32[9] { - %param_0.4541 = f32[9]{0:T(128)} parameter(0) - %param_1.5338 = s32[256]{0:T(256)} parameter(1) - %reshape.3919 = s32[256]{0:T(256)} reshape(%param_1.5338), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} - %transpose.1108 = s32[256]{0:T(256)} transpose(%reshape.3919), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} - %param_2.4503 = f32[256]{0:T(256)} parameter(2) - %reshape.3920 = f32[256]{0:T(256)} reshape(%param_2.4503), metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} - %transpose.1109 = f32[256]{0:T(256)} transpose(%reshape.3920), dimensions={0}, metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} - ROOT %scatter-add.235 = f32[9]{0:T(128)} scatter(%param_0.4541, %transpose.1108, %transpose.1109), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, indices_are_sorted=true, to_apply=%region_79.95.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} -}, execution_thread="sparsecore" - -%called_computation.22 (param_0.4542: f32[9], param_1.5339: s32[256], param_2.4504: f32[256]) -> f32[9] { - %param_0.4542 = f32[9]{0:T(128)} parameter(0) - %param_1.5339 = s32[256]{0:T(256)} parameter(1) - %param_2.4504 = f32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.47 = f32[9]{0:T(128)} fusion(%param_0.4542, %param_1.5339, %param_2.4504), kind=kCustom, calls=%fused_computation.27.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["16"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"1312","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} -}, execution_thread="sparsecore" - -%async_computation.22 (param_0.4543: f32[9], param_1.5340: s32[256], param_2.4505: f32[256]) -> f32[9] { - %param_0.4543 = f32[9]{0:T(128)} parameter(0) - %param_1.5340 = s32[256]{0:T(256)} parameter(1) - %param_2.4505 = f32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.48.cloned.1 = f32[9]{0:T(128)} call(%param_0.4543, %param_1.5340, %param_2.4505), to_apply=%called_computation.22, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} -}, execution_thread="sparsecore" - -%called_computation.4 (param_0.96: f32[9], param_1.148: s32[256], param_2.92: f32[256], param_3.3097: token[]) -> f32[9] { - %param_3.3097 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} - %param_0.96 = f32[9]{0:T(128)} parameter(0), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} - %param_1.148 = s32[256]{0:T(256)} parameter(1), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} - %param_2.92 = f32[256]{0:T(256)} parameter(2), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} - %copy.2082.cloned.1.call-start = ((f32[9]{0:T(128)}), f32[9]{0:T(128)}, u32[]{:S(8)}) async-start(%param_0.96), async_execution_thread="sparsecore", calls=%async_computation.21 - %copy.2082.cloned.1.call-done = f32[9]{0:T(128)} async-done(%copy.2082.cloned.1.call-start) - %scatter_offload_custom_fusion.48.cloned.1.call-start = ((f32[9]{0:T(128)}, s32[256]{0:T(256)}, f32[256]{0:T(256)}), f32[9]{0:T(128)}, u32[]{:S(8)}) async-start(%copy.2082.cloned.1.call-done, %param_1.148, %param_2.92), async_execution_thread="sparsecore", calls=%async_computation.22, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} - ROOT %scatter_offload_custom_fusion.48.cloned.1.call-done = f32[9]{0:T(128)} async-done(%scatter_offload_custom_fusion.48.cloned.1.call-start), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} -}, execution_thread="sparsecore" - -%async_computation.4 (param_0.97: f32[9], param_1.149: s32[256], param_2.93: f32[256], param_3.3096: token[]) -> f32[9] { - %param_3.3096 = token[] parameter(3) - %param_0.97 = f32[9]{0:T(128)} parameter(0) - %param_1.149 = s32[256]{0:T(256)} parameter(1) - %param_2.93 = f32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.14.cloned.1 = f32[9]{0:T(128)} call(%param_0.97, %param_1.149, %param_2.93, %param_3.3096), to_apply=%called_computation.4, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} -}, execution_thread="sparsecore" - -%called_computation.23 (param_0.4544: s32[263]) -> s32[263] { - %param_0.4544 = s32[263]{0:T(512)} parameter(0) - ROOT %copy.2083 = s32[263]{0:T(512)} copy(%param_0.4544), backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["32"],"input_window_bounds":[],"estimated_cycles":"1141","iteration_bounds":[],"scratchpad_allocation_size":"512","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"16","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} -}, execution_thread="sparsecore" - -%async_computation.23 (param_0.4545: s32[263]) -> s32[263] { - %param_0.4545 = s32[263]{0:T(512)} parameter(0) - ROOT %copy.2084.cloned.1 = s32[263]{0:T(512)} call(%param_0.4545), to_apply=%called_computation.23 -}, execution_thread="sparsecore" - -%region_81.97.clone (scatter-add.171: s32[], scatter-add.172: s32[]) -> s32[] { - %scatter-add.171 = s32[]{:T(128)S(7)} parameter(0), metadata={op_name="scatter-add"} - %scatter-add.172 = s32[]{:T(128)S(7)} parameter(1), metadata={op_name="scatter-add"} - ROOT %add.2531 = s32[]{:T(128)S(7)} add(%scatter-add.171, %scatter-add.172), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} -}, execution_thread="sparsecore" - -%fused_computation.28.clone.clone (param_0.4546: s32[263], param_1.5341: s32[8], param_2.4506: s32[8]) -> s32[263] { - %param_0.4546 = s32[263]{0:T(512)} parameter(0) - %param_1.5341 = s32[8]{0:T(128)} parameter(1) - %reshape.3921 = s32[8]{0:T(128)} reshape(%param_1.5341), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/select_n" stack_frame_id=0} - %transpose.1110 = s32[8]{0:T(128)} transpose(%reshape.3921), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/select_n" stack_frame_id=0} - %param_2.4506 = s32[8]{0:T(128)} parameter(2) - %reshape.3922 = s32[8]{0:T(128)} reshape(%param_2.4506), metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast.80" stack_frame_id=0} - %transpose.1111 = s32[8]{0:T(128)} transpose(%reshape.3922), dimensions={0}, metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast.80" stack_frame_id=0} - ROOT %scatter-add.236 = s32[263]{0:T(512)} scatter(%param_0.4546, %transpose.1110, %transpose.1111), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, to_apply=%region_81.97.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} -}, execution_thread="sparsecore" - -%called_computation.24 (param_0.4547: s32[263], param_1.5342: s32[8], param_2.4507: s32[8]) -> s32[263] { - %param_0.4547 = s32[263]{0:T(512)} parameter(0) - %param_1.5342 = s32[8]{0:T(128)} parameter(1) - %param_2.4507 = s32[8]{0:T(128)} parameter(2) - ROOT %scatter_offload_custom_fusion.49 = s32[263]{0:T(512)} fusion(%param_0.4547, %param_1.5342, %param_2.4507), kind=kCustom, calls=%fused_computation.28.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["8"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"256","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} -}, execution_thread="sparsecore" - -%async_computation.24 (param_0.4548: s32[263], param_1.5343: s32[8], param_2.4508: s32[8]) -> s32[263] { - %param_0.4548 = s32[263]{0:T(512)} parameter(0) - %param_1.5343 = s32[8]{0:T(128)} parameter(1) - %param_2.4508 = s32[8]{0:T(128)} parameter(2) - ROOT %scatter_offload_custom_fusion.50.cloned.1 = s32[263]{0:T(512)} call(%param_0.4548, %param_1.5343, %param_2.4508), to_apply=%called_computation.24, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} -}, execution_thread="sparsecore" - -%called_computation.5 (param_0.99: s32[263], param_1.151: s32[8], param_2.95: s32[8], param_3.3107: token[]) -> s32[263] { - %param_3.3107 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} - %param_0.99 = s32[263]{0:T(512)} parameter(0), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} - %param_1.151 = s32[8]{0:T(128)} parameter(1), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} - %param_2.95 = s32[8]{0:T(128)} parameter(2), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} - %copy.2084.cloned.1.call-start = ((s32[263]{0:T(512)}), s32[263]{0:T(512)}, u32[]{:S(8)}) async-start(%param_0.99), async_execution_thread="sparsecore", calls=%async_computation.23 - %copy.2084.cloned.1.call-done = s32[263]{0:T(512)} async-done(%copy.2084.cloned.1.call-start) - %scatter_offload_custom_fusion.50.cloned.1.call-start = ((s32[263]{0:T(512)}, s32[8]{0:T(128)}, s32[8]{0:T(128)}), s32[263]{0:T(512)}, u32[]{:S(8)}) async-start(%copy.2084.cloned.1.call-done, %param_1.151, %param_2.95), async_execution_thread="sparsecore", calls=%async_computation.24, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} - ROOT %scatter_offload_custom_fusion.50.cloned.1.call-done = s32[263]{0:T(512)} async-done(%scatter_offload_custom_fusion.50.cloned.1.call-start), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} -}, execution_thread="sparsecore" - -%async_computation.5 (param_0.100: s32[263], param_1.152: s32[8], param_2.96: s32[8], param_3.3106: token[]) -> s32[263] { - %param_3.3106 = token[] parameter(3) - %param_0.100 = s32[263]{0:T(512)} parameter(0) - %param_1.152 = s32[8]{0:T(128)} parameter(1) - %param_2.96 = s32[8]{0:T(128)} parameter(2) - ROOT %scatter_offload_custom_fusion.17.cloned.1 = s32[263]{0:T(512)} call(%param_0.100, %param_1.152, %param_2.96, %param_3.3106), to_apply=%called_computation.5, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} -}, execution_thread="sparsecore" - -%called_computation.25 (param_0.4549: s32[263]) -> s32[263] { - %param_0.4549 = s32[263]{0:T(512)} parameter(0) - ROOT %copy.2085 = s32[263]{0:T(512)} copy(%param_0.4549), backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["32"],"input_window_bounds":[],"estimated_cycles":"1141","iteration_bounds":[],"scratchpad_allocation_size":"512","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"16","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} -}, execution_thread="sparsecore" - -%async_computation.25 (param_0.4550: s32[263]) -> s32[263] { - %param_0.4550 = s32[263]{0:T(512)} parameter(0) - ROOT %copy.2086.cloned.1 = s32[263]{0:T(512)} call(%param_0.4550), to_apply=%called_computation.25 -}, execution_thread="sparsecore" - -%region_96.114 (scatter-add.48: s32[], scatter-add.49: s32[]) -> s32[] { - %scatter-add.48 = s32[]{:T(128)S(7)} parameter(0), metadata={op_name="scatter-add"} - %scatter-add.49 = s32[]{:T(128)S(7)} parameter(1), metadata={op_name="scatter-add"} - ROOT %add.1434 = s32[]{:T(128)S(7)} add(%scatter-add.48, %scatter-add.49), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} -}, execution_thread="sparsecore" - -%fused_computation.29.clone.clone (param_0.4551: s32[263], param_1.5344: s32[256], param_2.4509: s32[256]) -> s32[263] { - %param_0.4551 = s32[263]{0:T(512)} parameter(0) - %param_1.5344 = s32[256]{0:T(256)} parameter(1) - %reshape.3923 = s32[256]{0:T(256)} reshape(%param_1.5344), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/select_n" stack_frame_id=0} - %transpose.1112 = s32[256]{0:T(256)} transpose(%reshape.3923), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/select_n" stack_frame_id=0} - %param_2.4509 = s32[256]{0:T(256)} parameter(2) - %reshape.3924 = s32[256]{0:T(256)} reshape(%param_2.4509), metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} - %transpose.1113 = s32[256]{0:T(256)} transpose(%reshape.3924), dimensions={0}, metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} - ROOT %scatter-add.237 = s32[263]{0:T(512)} scatter(%param_0.4551, %transpose.1112, %transpose.1113), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, to_apply=%region_96.114, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} -}, execution_thread="sparsecore" - -%called_computation.26 (param_0.4552: s32[263], param_1.5345: s32[256], param_2.4510: s32[256]) -> s32[263] { - %param_0.4552 = s32[263]{0:T(512)} parameter(0) - %param_1.5345 = s32[256]{0:T(256)} parameter(1) - %param_2.4510 = s32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.51 = s32[263]{0:T(512)} fusion(%param_0.4552, %param_1.5345, %param_2.4510), kind=kCustom, calls=%fused_computation.29.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["16"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"384","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} -}, execution_thread="sparsecore" - -%async_computation.26 (param_0.4553: s32[263], param_1.5346: s32[256], param_2.4511: s32[256]) -> s32[263] { - %param_0.4553 = s32[263]{0:T(512)} parameter(0) - %param_1.5346 = s32[256]{0:T(256)} parameter(1) - %param_2.4511 = s32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.52.cloned.1 = s32[263]{0:T(512)} call(%param_0.4553, %param_1.5346, %param_2.4511), to_apply=%called_computation.26, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} -}, execution_thread="sparsecore" - -%called_computation.6 (param_0.102: s32[263], param_1.154: s32[256], param_2.98: s32[256], param_3.3093: token[]) -> s32[263] { - %param_3.3093 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} - %param_0.102 = s32[263]{0:T(512)} parameter(0), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} - %param_1.154 = s32[256]{0:T(256)} parameter(1), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} - %param_2.98 = s32[256]{0:T(256)} parameter(2), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} - %copy.2086.cloned.1.call-start = ((s32[263]{0:T(512)}), s32[263]{0:T(512)}, u32[]{:S(8)}) async-start(%param_0.102), async_execution_thread="sparsecore", calls=%async_computation.25 - %copy.2086.cloned.1.call-done = s32[263]{0:T(512)} async-done(%copy.2086.cloned.1.call-start) - %scatter_offload_custom_fusion.52.cloned.1.call-start = ((s32[263]{0:T(512)}, s32[256]{0:T(256)}, s32[256]{0:T(256)}), s32[263]{0:T(512)}, u32[]{:S(8)}) async-start(%copy.2086.cloned.1.call-done, %param_1.154, %param_2.98), async_execution_thread="sparsecore", calls=%async_computation.26, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} - ROOT %scatter_offload_custom_fusion.52.cloned.1.call-done = s32[263]{0:T(512)} async-done(%scatter_offload_custom_fusion.52.cloned.1.call-start), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} -}, execution_thread="sparsecore" - -%async_computation.6 (param_0.103: s32[263], param_1.155: s32[256], param_2.99: s32[256], param_3.3092: token[]) -> s32[263] { - %param_3.3092 = token[] parameter(3) - %param_0.103 = s32[263]{0:T(512)} parameter(0) - %param_1.155 = s32[256]{0:T(256)} parameter(1) - %param_2.99 = s32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.20.cloned.1 = s32[263]{0:T(512)} call(%param_0.103, %param_1.155, %param_2.99, %param_3.3092), to_apply=%called_computation.6, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} -}, execution_thread="sparsecore" - -%region_102.120 (scatter-add.52: f32[], scatter-add.53: f32[]) -> f32[] { - %scatter-add.52 = f32[]{:T(128)S(7)} parameter(0), metadata={op_name="scatter-add"} - %scatter-add.53 = f32[]{:T(128)S(7)} parameter(1), metadata={op_name="scatter-add"} - ROOT %add.1437 = f32[]{:T(128)S(7)} add(%scatter-add.52, %scatter-add.53), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} -}, execution_thread="sparsecore" - -%fused_computation.30.clone.clone (param_0.4556: f32[9], param_1.5347: s32[256], param_2.4512: f32[256]) -> f32[9] { - %param_0.4556 = f32[9]{0:T(128)} parameter(0) - %param_1.5347 = s32[256]{0:T(256)} parameter(1) - %reshape.3925 = s32[256]{0:T(256)} reshape(%param_1.5347), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/broadcast_in_dim" stack_frame_id=0} - %transpose.1114 = s32[256]{0:T(256)} transpose(%reshape.3925), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/broadcast_in_dim" stack_frame_id=0} - %param_2.4512 = f32[256]{0:T(256)} parameter(2) - %reshape.3926 = f32[256]{0:T(256)} reshape(%param_2.4512), metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} - %transpose.1115 = f32[256]{0:T(256)} transpose(%reshape.3926), dimensions={0}, metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} - ROOT %scatter-add.238 = f32[9]{0:T(128)} scatter(%param_0.4556, %transpose.1114, %transpose.1115), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, indices_are_sorted=true, to_apply=%region_102.120, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} -}, execution_thread="sparsecore" - -%called_computation.28 (param_0.4557: f32[9], param_1.5348: s32[256], param_2.4513: f32[256]) -> f32[9] { - %param_0.4557 = f32[9]{0:T(128)} parameter(0) - %param_1.5348 = s32[256]{0:T(256)} parameter(1) - %param_2.4513 = f32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.53 = f32[9]{0:T(128)} fusion(%param_0.4557, %param_1.5348, %param_2.4513), kind=kCustom, calls=%fused_computation.30.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["16"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"1312","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} -}, execution_thread="sparsecore" - -%async_computation.28 (param_0.4558: f32[9], param_1.5349: s32[256], param_2.4514: f32[256]) -> f32[9] { - %param_0.4558 = f32[9]{0:T(128)} parameter(0) - %param_1.5349 = s32[256]{0:T(256)} parameter(1) - %param_2.4514 = f32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.54.cloned.1 = f32[9]{0:T(128)} call(%param_0.4558, %param_1.5349, %param_2.4514), to_apply=%called_computation.28, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} -}, execution_thread="sparsecore" - -%called_computation.7 (param_0.105: f32[9], param_1.157: s32[256], param_2.101: f32[256], param_3.3101: token[]) -> f32[9] { - %param_3.3101 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} - %param_0.105 = f32[9]{0:T(128)} parameter(0), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} - %param_1.157 = s32[256]{0:T(256)} parameter(1), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} - %param_2.101 = f32[256]{0:T(256)} parameter(2), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} - %scatter_offload_custom_fusion.54.cloned.1.call-start = ((f32[9]{0:T(128)}, s32[256]{0:T(256)}, f32[256]{0:T(256)}), f32[9]{0:T(128)}, u32[]{:S(8)}) async-start(%param_0.105, %param_1.157, %param_2.101), async_execution_thread="sparsecore", calls=%async_computation.28, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} - ROOT %scatter_offload_custom_fusion.54.cloned.1.call-done = f32[9]{0:T(128)} async-done(%scatter_offload_custom_fusion.54.cloned.1.call-start), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} -}, execution_thread="sparsecore" - -%async_computation.7 (param_0.106: f32[9], param_1.158: s32[256], param_2.102: f32[256], param_3.3100: token[]) -> f32[9] { - %param_3.3100 = token[] parameter(3) - %param_0.106 = f32[9]{0:T(128)} parameter(0) - %param_1.158 = s32[256]{0:T(256)} parameter(1) - %param_2.102 = f32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.23.cloned.1 = f32[9]{0:T(128)} call(%param_0.106, %param_1.158, %param_2.102, %param_3.3100), to_apply=%called_computation.7, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} -}, execution_thread="sparsecore" - -%region_104.122 (scatter-add.83: s32[], scatter-add.84: s32[]) -> s32[] { - %scatter-add.83 = s32[]{:T(128)S(7)} parameter(0), metadata={op_name="scatter-add"} - %scatter-add.84 = s32[]{:T(128)S(7)} parameter(1), metadata={op_name="scatter-add"} - ROOT %add.1438 = s32[]{:T(128)S(7)} add(%scatter-add.83, %scatter-add.84), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} -}, execution_thread="sparsecore" - -%fused_computation.31.clone.clone (param_0.4561: s32[263], param_1.5350: s32[8], param_2.4515: s32[8]) -> s32[263] { - %param_0.4561 = s32[263]{0:T(512)} parameter(0) - %param_1.5350 = s32[8]{0:T(128)} parameter(1) - %reshape.3927 = s32[8]{0:T(128)} reshape(%param_1.5350), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/select_n" stack_frame_id=0} - %transpose.1116 = s32[8]{0:T(128)} transpose(%reshape.3927), dimensions={0}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/select_n" stack_frame_id=0} - %param_2.4515 = s32[8]{0:T(128)} parameter(2) - %reshape.3928 = s32[8]{0:T(128)} reshape(%param_2.4515), metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast.80" stack_frame_id=0} - %transpose.1117 = s32[8]{0:T(128)} transpose(%reshape.3928), dimensions={0}, metadata={op_name="jit(train_step)/moe_layers/shard_map/jit(gmm)/broadcast.80" stack_frame_id=0} - ROOT %scatter-add.239 = s32[263]{0:T(512)} scatter(%param_0.4561, %transpose.1116, %transpose.1117), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, to_apply=%region_104.122, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} -}, execution_thread="sparsecore" - -%called_computation.30 (param_0.4562: s32[263], param_1.5351: s32[8], param_2.4516: s32[8]) -> s32[263] { - %param_0.4562 = s32[263]{0:T(512)} parameter(0) - %param_1.5351 = s32[8]{0:T(128)} parameter(1) - %param_2.4516 = s32[8]{0:T(128)} parameter(2) - ROOT %scatter_offload_custom_fusion.55 = s32[263]{0:T(512)} fusion(%param_0.4562, %param_1.5351, %param_2.4516), kind=kCustom, calls=%fused_computation.31.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["8"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"256","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} -}, execution_thread="sparsecore" - -%async_computation.30 (param_0.4563: s32[263], param_1.5352: s32[8], param_2.4517: s32[8]) -> s32[263] { - %param_0.4563 = s32[263]{0:T(512)} parameter(0) - %param_1.5352 = s32[8]{0:T(128)} parameter(1) - %param_2.4517 = s32[8]{0:T(128)} parameter(2) - ROOT %scatter_offload_custom_fusion.56.cloned.1 = s32[263]{0:T(512)} call(%param_0.4563, %param_1.5352, %param_2.4517), to_apply=%called_computation.30, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} -}, execution_thread="sparsecore" - -%called_computation.8 (param_0.108: s32[263], param_1.160: s32[8], param_2.104: s32[8], param_3.3109: token[]) -> s32[263] { - %param_3.3109 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} - %param_0.108 = s32[263]{0:T(512)} parameter(0), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} - %param_1.160 = s32[8]{0:T(128)} parameter(1), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} - %param_2.104 = s32[8]{0:T(128)} parameter(2), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} - %scatter_offload_custom_fusion.56.cloned.1.call-start = ((s32[263]{0:T(512)}, s32[8]{0:T(128)}, s32[8]{0:T(128)}), s32[263]{0:T(512)}, u32[]{:S(8)}) async-start(%param_0.108, %param_1.160, %param_2.104), async_execution_thread="sparsecore", calls=%async_computation.30, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} - ROOT %scatter_offload_custom_fusion.56.cloned.1.call-done = s32[263]{0:T(512)} async-done(%scatter_offload_custom_fusion.56.cloned.1.call-start), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} -}, execution_thread="sparsecore" - -%async_computation.8 (param_0.109: s32[263], param_1.161: s32[8], param_2.105: s32[8], param_3.3108: token[]) -> s32[263] { - %param_3.3108 = token[] parameter(3) - %param_0.109 = s32[263]{0:T(512)} parameter(0) - %param_1.161 = s32[8]{0:T(128)} parameter(1) - %param_2.105 = s32[8]{0:T(128)} parameter(2) - ROOT %scatter_offload_custom_fusion.26.cloned.1 = s32[263]{0:T(512)} call(%param_0.109, %param_1.161, %param_2.105, %param_3.3108), to_apply=%called_computation.8, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/shard_map/jit(tgmm)/scatter-add" stack_frame_id=0} -}, execution_thread="sparsecore" - -%region_14.20 (scatter-add.0: s32[], scatter-add.1: s32[]) -> s32[] { - %scatter-add.0 = s32[]{:T(128)S(7)} parameter(0), metadata={op_name="scatter-add"} - %scatter-add.1 = s32[]{:T(128)S(7)} parameter(1), metadata={op_name="scatter-add"} - ROOT %add.1345 = s32[]{:T(128)S(7)} add(%scatter-add.0, %scatter-add.1), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} -}, execution_thread="sparsecore" - -%fused_computation.17.clone.clone.clone (param_0.4566: s32[256], param_1.5353: s32[4096], param_2.4518: s32[4096]) -> s32[256] { - %param_0.4566 = s32[256]{0:T(256)} parameter(0) - %param_1.5353 = s32[4096]{0:T(1024)} parameter(1) - %reshape.3929 = s32[4096]{0:T(1024)} reshape(%param_1.5353), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/select_n" stack_frame_id=0} - %transpose.1118 = s32[4096]{0:T(1024)} transpose(%reshape.3929), dimensions={0}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/select_n" stack_frame_id=0} - %param_2.4518 = s32[4096]{0:T(1024)} parameter(2) - %reshape.3930 = s32[4096]{0:T(1024)} reshape(%param_2.4518), metadata={op_name="jit(train_step)/moe_layers/shard_map/broadcast_in_dim" stack_frame_id=0} - %transpose.1119 = s32[4096]{0:T(1024)} transpose(%reshape.3930), dimensions={0}, metadata={op_name="jit(train_step)/moe_layers/shard_map/broadcast_in_dim" stack_frame_id=0} - ROOT %scatter-add.240 = s32[256]{0:T(256)} scatter(%param_0.4566, %transpose.1118, %transpose.1119), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, to_apply=%region_14.20, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/scatter-add" stack_frame_id=0} -}, execution_thread="sparsecore" - -%called_computation.32 (param_0.4567: s32[256], param_1.5354: s32[4096], param_2.4519: s32[4096]) -> s32[256] { - %param_0.4567 = s32[256]{0:T(256)} parameter(0) - %param_1.5354 = s32[4096]{0:T(1024)} parameter(1) - %param_2.4519 = s32[4096]{0:T(1024)} parameter(2) - ROOT %scatter_offload_custom_fusion.57 = s32[256]{0:T(256)} fusion(%param_0.4567, %param_1.5354, %param_2.4519), kind=kCustom, calls=%fused_computation.17.clone.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["256"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"4160","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} -}, execution_thread="sparsecore" - -%async_computation.32 (param_0.4568: s32[256], param_1.5355: s32[4096], param_2.4520: s32[4096]) -> s32[256] { - %param_0.4568 = s32[256]{0:T(256)} parameter(0) - %param_1.5355 = s32[4096]{0:T(1024)} parameter(1) - %param_2.4520 = s32[4096]{0:T(1024)} parameter(2) - ROOT %scatter_offload_custom_fusion.58.cloned.1 = s32[256]{0:T(256)} call(%param_0.4568, %param_1.5355, %param_2.4520), to_apply=%called_computation.32, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/scatter-add" stack_frame_id=0} -}, execution_thread="sparsecore" - -%called_computation.9 (param_0.111: s32[256], param_1.163: s32[4096], param_2.107: s32[4096], param_3.3087: token[]) -> s32[256] { - %param_3.3087 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} - %param_0.111 = s32[256]{0:T(256)} parameter(0), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} - %param_1.163 = s32[4096]{0:T(1024)} parameter(1), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} - %param_2.107 = s32[4096]{0:T(1024)} parameter(2), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} - %scatter_offload_custom_fusion.58.cloned.1.call-start = ((s32[256]{0:T(256)}, s32[4096]{0:T(1024)}, s32[4096]{0:T(1024)}), s32[256]{0:T(256)}, u32[]{:S(8)}) async-start(%param_0.111, %param_1.163, %param_2.107), async_execution_thread="sparsecore", calls=%async_computation.32, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/scatter-add" stack_frame_id=0} - ROOT %scatter_offload_custom_fusion.58.cloned.1.call-done = s32[256]{0:T(256)} async-done(%scatter_offload_custom_fusion.58.cloned.1.call-start), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/scatter-add" stack_frame_id=0} -}, execution_thread="sparsecore" - -%async_computation.9 (param_0.112: s32[256], param_1.164: s32[4096], param_2.108: s32[4096], param_3.3086: token[]) -> s32[256] { - %param_3.3086 = token[] parameter(3) - %param_0.112 = s32[256]{0:T(256)} parameter(0) - %param_1.164 = s32[4096]{0:T(1024)} parameter(1) - %param_2.108 = s32[4096]{0:T(1024)} parameter(2) - ROOT %scatter_offload_custom_fusion.29.cloned.1 = s32[256]{0:T(256)} call(%param_0.112, %param_1.164, %param_2.108, %param_3.3086), to_apply=%called_computation.9, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/scatter-add" stack_frame_id=0} -}, execution_thread="sparsecore" - -%called_computation.33 (param_0.4569: s32[263]) -> s32[263] { - %param_0.4569 = s32[263]{0:T(512)} parameter(0) - ROOT %copy.2093 = s32[263]{0:T(512)} copy(%param_0.4569), backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["32"],"input_window_bounds":[],"estimated_cycles":"1141","iteration_bounds":[],"scratchpad_allocation_size":"512","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"16","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} -}, execution_thread="sparsecore" - -%async_computation.33 (param_0.4570: s32[263]) -> s32[263] { - %param_0.4570 = s32[263]{0:T(512)} parameter(0) - ROOT %copy.2094.cloned.1 = s32[263]{0:T(512)} call(%param_0.4570), to_apply=%called_computation.33 -}, execution_thread="sparsecore" - -%region_20.26.clone.1 (scatter-add.141: s32[], scatter-add.142: s32[]) -> s32[] { - %scatter-add.141 = s32[]{:T(128)S(7)} parameter(0), metadata={op_name="scatter-add"} - %scatter-add.142 = s32[]{:T(128)S(7)} parameter(1), metadata={op_name="scatter-add"} - ROOT %add.2516 = s32[]{:T(128)S(7)} add(%scatter-add.141, %scatter-add.142), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} -}, execution_thread="sparsecore" - -%fused_computation.18.clone.clone.clone (param_0.4571: s32[263], param_1.5356: s32[256], param_2.4521: s32[256]) -> s32[263] { - %param_0.4571 = s32[263]{0:T(512)} parameter(0) - %param_1.5356 = s32[256]{0:T(256)} parameter(1) - %reshape.3931 = s32[256]{0:T(256)} reshape(%param_1.5356), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/select_n" stack_frame_id=0} - %transpose.1120 = s32[256]{0:T(256)} transpose(%reshape.3931), dimensions={0}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/select_n" stack_frame_id=0} - %param_2.4521 = s32[256]{0:T(256)} parameter(2) - %reshape.3932 = s32[256]{0:T(256)} reshape(%param_2.4521) - %transpose.1121 = s32[256]{0:T(256)} transpose(%reshape.3932), dimensions={0} - ROOT %scatter-add.241 = s32[263]{0:T(512)} scatter(%param_0.4571, %transpose.1120, %transpose.1121), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, to_apply=%region_20.26.clone.1, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} -}, execution_thread="sparsecore" - -%called_computation.34 (param_0.4572: s32[263], param_1.5357: s32[256], param_2.4522: s32[256]) -> s32[263] { - %param_0.4572 = s32[263]{0:T(512)} parameter(0) - %param_1.5357 = s32[256]{0:T(256)} parameter(1) - %param_2.4522 = s32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.59 = s32[263]{0:T(512)} fusion(%param_0.4572, %param_1.5357, %param_2.4522), kind=kCustom, calls=%fused_computation.18.clone.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["16"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"384","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} -}, execution_thread="sparsecore" - -%async_computation.34 (param_0.4573: s32[263], param_1.5358: s32[256], param_2.4523: s32[256]) -> s32[263] { - %param_0.4573 = s32[263]{0:T(512)} parameter(0) - %param_1.5358 = s32[256]{0:T(256)} parameter(1) - %param_2.4523 = s32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.60.cloned.1 = s32[263]{0:T(512)} call(%param_0.4573, %param_1.5358, %param_2.4523), to_apply=%called_computation.34, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} -}, execution_thread="sparsecore" - -%called_computation.10 (param_0.114: s32[263], param_1.166: s32[256], param_2.110: s32[256], param_3.3089: token[]) -> s32[263] { - %param_3.3089 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} - %param_0.114 = s32[263]{0:T(512)} parameter(0), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} - %param_1.166 = s32[256]{0:T(256)} parameter(1), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} - %param_2.110 = s32[256]{0:T(256)} parameter(2), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} - %copy.2094.cloned.1.call-start = ((s32[263]{0:T(512)}), s32[263]{0:T(512)}, u32[]{:S(8)}) async-start(%param_0.114), async_execution_thread="sparsecore", calls=%async_computation.33 - %copy.2094.cloned.1.call-done = s32[263]{0:T(512)} async-done(%copy.2094.cloned.1.call-start) - %scatter_offload_custom_fusion.60.cloned.1.call-start = ((s32[263]{0:T(512)}, s32[256]{0:T(256)}, s32[256]{0:T(256)}), s32[263]{0:T(512)}, u32[]{:S(8)}) async-start(%copy.2094.cloned.1.call-done, %param_1.166, %param_2.110), async_execution_thread="sparsecore", calls=%async_computation.34, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} - ROOT %scatter_offload_custom_fusion.60.cloned.1.call-done = s32[263]{0:T(512)} async-done(%scatter_offload_custom_fusion.60.cloned.1.call-start), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} -}, execution_thread="sparsecore" - -%async_computation.10 (param_0.115: s32[263], param_1.167: s32[256], param_2.111: s32[256], param_3.3088: token[]) -> s32[263] { - %param_3.3088 = token[] parameter(3) - %param_0.115 = s32[263]{0:T(512)} parameter(0) - %param_1.167 = s32[256]{0:T(256)} parameter(1) - %param_2.111 = s32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.32.cloned.1 = s32[263]{0:T(512)} call(%param_0.115, %param_1.167, %param_2.111, %param_3.3088), to_apply=%called_computation.10, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} -}, execution_thread="sparsecore" - -%called_computation.35 (param_0.4574: f32[9]) -> f32[9] { - %param_0.4574 = f32[9]{0:T(128)} parameter(0) - ROOT %copy.2095 = f32[9]{0:T(128)} copy(%param_0.4574), backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["8"],"input_window_bounds":[],"estimated_cycles":"1131","iteration_bounds":[],"scratchpad_allocation_size":"128","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"16","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} -}, execution_thread="sparsecore" - -%async_computation.35 (param_0.4575: f32[9]) -> f32[9] { - %param_0.4575 = f32[9]{0:T(128)} parameter(0) - ROOT %copy.2096.cloned.1 = f32[9]{0:T(128)} call(%param_0.4575), to_apply=%called_computation.35 -}, execution_thread="sparsecore" - -%region_26.33.clone.1 (scatter-add.145: f32[], scatter-add.146: f32[]) -> f32[] { - %scatter-add.145 = f32[]{:T(128)S(7)} parameter(0), metadata={op_name="scatter-add"} - %scatter-add.146 = f32[]{:T(128)S(7)} parameter(1), metadata={op_name="scatter-add"} - ROOT %add.2518 = f32[]{:T(128)S(7)} add(%scatter-add.145, %scatter-add.146), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} -}, execution_thread="sparsecore" - -%fused_computation.19.clone.clone.clone (param_0.4576: f32[9], param_1.5359: s32[256], param_2.4524: f32[256]) -> f32[9] { - %param_0.4576 = f32[9]{0:T(128)} parameter(0) - %param_1.5359 = s32[256]{0:T(256)} parameter(1) - %reshape.3933 = s32[256]{0:T(256)} reshape(%param_1.5359), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} - %transpose.1122 = s32[256]{0:T(256)} transpose(%reshape.3933), dimensions={0}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/broadcast_in_dim" stack_frame_id=0} - %param_2.4524 = f32[256]{0:T(256)} parameter(2) - %reshape.3934 = f32[256]{0:T(256)} reshape(%param_2.4524) - %transpose.1123 = f32[256]{0:T(256)} transpose(%reshape.3934), dimensions={0} - ROOT %scatter-add.242 = f32[9]{0:T(128)} scatter(%param_0.4576, %transpose.1122, %transpose.1123), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, indices_are_sorted=true, to_apply=%region_26.33.clone.1, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} -}, execution_thread="sparsecore" - -%called_computation.36 (param_0.4577: f32[9], param_1.5360: s32[256], param_2.4525: f32[256]) -> f32[9] { - %param_0.4577 = f32[9]{0:T(128)} parameter(0) - %param_1.5360 = s32[256]{0:T(256)} parameter(1) - %param_2.4525 = f32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.61 = f32[9]{0:T(128)} fusion(%param_0.4577, %param_1.5360, %param_2.4525), kind=kCustom, calls=%fused_computation.19.clone.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["16"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"1312","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} -}, execution_thread="sparsecore" - -%async_computation.36 (param_0.4578: f32[9], param_1.5361: s32[256], param_2.4526: f32[256]) -> f32[9] { - %param_0.4578 = f32[9]{0:T(128)} parameter(0) - %param_1.5361 = s32[256]{0:T(256)} parameter(1) - %param_2.4526 = f32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.62.cloned.1 = f32[9]{0:T(128)} call(%param_0.4578, %param_1.5361, %param_2.4526), to_apply=%called_computation.36, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} -}, execution_thread="sparsecore" - -%called_computation.11 (param_0.117: f32[9], param_1.169: s32[256], param_2.113: f32[256], param_3.3095: token[]) -> f32[9] { - %param_3.3095 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} - %param_0.117 = f32[9]{0:T(128)} parameter(0), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} - %param_1.169 = s32[256]{0:T(256)} parameter(1), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} - %param_2.113 = f32[256]{0:T(256)} parameter(2), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} - %copy.2096.cloned.1.call-start = ((f32[9]{0:T(128)}), f32[9]{0:T(128)}, u32[]{:S(8)}) async-start(%param_0.117), async_execution_thread="sparsecore", calls=%async_computation.35 - %copy.2096.cloned.1.call-done = f32[9]{0:T(128)} async-done(%copy.2096.cloned.1.call-start) - %scatter_offload_custom_fusion.62.cloned.1.call-start = ((f32[9]{0:T(128)}, s32[256]{0:T(256)}, f32[256]{0:T(256)}), f32[9]{0:T(128)}, u32[]{:S(8)}) async-start(%copy.2096.cloned.1.call-done, %param_1.169, %param_2.113), async_execution_thread="sparsecore", calls=%async_computation.36, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} - ROOT %scatter_offload_custom_fusion.62.cloned.1.call-done = f32[9]{0:T(128)} async-done(%scatter_offload_custom_fusion.62.cloned.1.call-start), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} -}, execution_thread="sparsecore" - -%async_computation.11 (param_0.118: f32[9], param_1.170: s32[256], param_2.114: f32[256], param_3.3094: token[]) -> f32[9] { - %param_3.3094 = token[] parameter(3) - %param_0.118 = f32[9]{0:T(128)} parameter(0) - %param_1.170 = s32[256]{0:T(256)} parameter(1) - %param_2.114 = f32[256]{0:T(256)} parameter(2) - ROOT %scatter_offload_custom_fusion.35.cloned.1 = f32[9]{0:T(128)} call(%param_0.118, %param_1.170, %param_2.114, %param_3.3094), to_apply=%called_computation.11, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} -}, execution_thread="sparsecore" - -%called_computation.37 (param_0.4579: s32[263]) -> s32[263] { - %param_0.4579 = s32[263]{0:T(512)} parameter(0) - ROOT %copy.2097 = s32[263]{0:T(512)} copy(%param_0.4579), backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["32"],"input_window_bounds":[],"estimated_cycles":"1141","iteration_bounds":[],"scratchpad_allocation_size":"512","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"16","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} -}, execution_thread="sparsecore" - -%async_computation.37 (param_0.4580: s32[263]) -> s32[263] { - %param_0.4580 = s32[263]{0:T(512)} parameter(0) - ROOT %copy.2098.cloned.1 = s32[263]{0:T(512)} call(%param_0.4580), to_apply=%called_computation.37 -}, execution_thread="sparsecore" - -%region_28.35.clone.1 (scatter-add.149: s32[], scatter-add.150: s32[]) -> s32[] { - %scatter-add.149 = s32[]{:T(128)S(7)} parameter(0), metadata={op_name="scatter-add"} - %scatter-add.150 = s32[]{:T(128)S(7)} parameter(1), metadata={op_name="scatter-add"} - ROOT %add.2520 = s32[]{:T(128)S(7)} add(%scatter-add.149, %scatter-add.150), metadata={op_name="add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["128"],"input_window_bounds":[],"estimated_cycles":"1165","iteration_bounds":[],"scratchpad_allocation_size":"520","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[{"unroll_dimension":"0","unroll_factor":"4","pipeline_remainder":false,"fully_unroll_if_trip_count_is_at_most":"0"}],"vectorizing_shape":[]},"scoped_memory_configs":[],"used_scoped_memory_configs":[]} -}, execution_thread="sparsecore" - -%fused_computation.20.clone.clone.clone (param_0.4581: s32[263], param_1.5362: s32[8], param_2.4527: s32[8]) -> s32[263] { - %param_0.4581 = s32[263]{0:T(512)} parameter(0) - %param_1.5362 = s32[8]{0:T(128)} parameter(1) - %reshape.3935 = s32[8]{0:T(128)} reshape(%param_1.5362), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/select_n" stack_frame_id=0} - %transpose.1124 = s32[8]{0:T(128)} transpose(%reshape.3935), dimensions={0}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/select_n" stack_frame_id=0} - %param_2.4527 = s32[8]{0:T(128)} parameter(2) - %reshape.3936 = s32[8]{0:T(128)} reshape(%param_2.4527) - %transpose.1125 = s32[8]{0:T(128)} transpose(%reshape.3936), dimensions={0} - ROOT %scatter-add.243 = s32[263]{0:T(512)} scatter(%param_0.4581, %transpose.1124, %transpose.1125), update_window_dims={}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=1, to_apply=%region_28.35.clone.1, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} -}, execution_thread="sparsecore" - -%called_computation.38 (param_0.4582: s32[263], param_1.5363: s32[8], param_2.4528: s32[8]) -> s32[263] { - %param_0.4582 = s32[263]{0:T(512)} parameter(0) - %param_1.5363 = s32[8]{0:T(128)} parameter(1) - %param_2.4528 = s32[8]{0:T(128)} parameter(2) - ROOT %scatter_offload_custom_fusion.63 = s32[263]{0:T(512)} fusion(%param_0.4582, %param_1.5363, %param_2.4528), kind=kCustom, calls=%fused_computation.20.clone.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0}, backend_config={"flag_configs":[],"window_config":{"kernel_window_bounds":[],"output_window_bounds":["8"],"input_window_bounds":[],"estimated_cycles":"9223372036854775807","iteration_bounds":[],"scratchpad_allocation_size":"256","cost_model_type":"COST_MODEL_TYPE_INVALID","ml_estimated_microseconds":0,"is_mask":false,"pad_output_on_minor_dim":"0","pad_input_on_minor_dim":"0","estimated_vmem_bytes":"0","estimated_bundle_count":"0","estimated_scoped_vmem_bytes":"0"},"loop_config":{"loop_order":[],"unrolled_loops":[],"vectorizing_shape":[]},"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_TILE","used_scoped_memory_configs":[]} -}, execution_thread="sparsecore" - -%async_computation.38 (param_0.4583: s32[263], param_1.5364: s32[8], param_2.4529: s32[8]) -> s32[263] { - %param_0.4583 = s32[263]{0:T(512)} parameter(0) - %param_1.5364 = s32[8]{0:T(128)} parameter(1) - %param_2.4529 = s32[8]{0:T(128)} parameter(2) - ROOT %scatter_offload_custom_fusion.64.cloned.1 = s32[263]{0:T(512)} call(%param_0.4583, %param_1.5364, %param_2.4529), to_apply=%called_computation.38, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} -}, execution_thread="sparsecore" - -%called_computation.12 (param_0.120: s32[263], param_1.172: s32[8], param_2.116: s32[8], param_3.3103: token[]) -> s32[263] { - %param_3.3103 = token[] parameter(3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"implicit_sharding":{"type":"REPLICATED","tile_assignment_dimensions":[],"tile_assignment_devices":[],"tuple_shardings":[],"replicate_on_last_tile_dim":false,"metadata":[],"last_tile_dims":[],"iota_reshape_dims":[],"iota_transpose_perm":[],"is_shard_group":false,"shard_group_id":"0","shard_group_type":"AS"},"used_scoped_memory_configs":[]} - %param_0.120 = s32[263]{0:T(512)} parameter(0), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} - %param_1.172 = s32[8]{0:T(128)} parameter(1), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} - %param_2.116 = s32[8]{0:T(128)} parameter(2), backend_config={"flag_configs":[],"scoped_memory_configs":[],"compute_type":"COMPUTE_TYPE_SCALAR","used_scoped_memory_configs":[]} - %copy.2098.cloned.1.call-start = ((s32[263]{0:T(512)}), s32[263]{0:T(512)}, u32[]{:S(8)}) async-start(%param_0.120), async_execution_thread="sparsecore", calls=%async_computation.37 - %copy.2098.cloned.1.call-done = s32[263]{0:T(512)} async-done(%copy.2098.cloned.1.call-start) - %scatter_offload_custom_fusion.64.cloned.1.call-start = ((s32[263]{0:T(512)}, s32[8]{0:T(128)}, s32[8]{0:T(128)}), s32[263]{0:T(512)}, u32[]{:S(8)}) async-start(%copy.2098.cloned.1.call-done, %param_1.172, %param_2.116), async_execution_thread="sparsecore", calls=%async_computation.38, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} - ROOT %scatter_offload_custom_fusion.64.cloned.1.call-done = s32[263]{0:T(512)} async-done(%scatter_offload_custom_fusion.64.cloned.1.call-start), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} -}, execution_thread="sparsecore" - -%async_computation.12 (param_0.121: s32[263], param_1.173: s32[8], param_2.117: s32[8], param_3.3102: token[]) -> s32[263] { - %param_3.3102 = token[] parameter(3) - %param_0.121 = s32[263]{0:T(512)} parameter(0) - %param_1.173 = s32[8]{0:T(128)} parameter(1) - %param_2.117 = s32[8]{0:T(128)} parameter(2) - ROOT %scatter_offload_custom_fusion.38.cloned.1 = s32[263]{0:T(512)} call(%param_0.121, %param_1.173, %param_2.117, %param_3.3102), to_apply=%called_computation.12, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/shard_map/jit(gmm)/scatter-add" stack_frame_id=0} -}, execution_thread="sparsecore" - -%region_154.179 (reduce_sum.431: f32[], reduce_sum.254: f32[]) -> f32[] { - %reduce_sum.431 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.254 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.258 = f32[]{:T(128)} add(%reduce_sum.431, %reduce_sum.254), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.466 (param_0.4166: f32[3,1536,128,192]) -> f32[] { - %param_0.4166 = f32[3,1536,128,192]{2,3,0,1:T(8,128)} parameter(0) - %bitcast.672 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} bitcast(%param_0.4166), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/dense_layers.wrapped_fn/transpose" stack_frame_id=0} - %square.564 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%bitcast.672, %bitcast.672), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.5101 = f32[]{:T(128)} constant(0) - ROOT %reduce.669 = f32[]{:T(128)} reduce(%square.564, %constant.5101), dimensions={0,1,2,3}, to_apply=%region_154.179, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} -} - -%fused_computation.467 (param_0.1420: f32[1536,3,128,192]) -> bf16[3,1536,128,192] { - %param_0.1420 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} parameter(0) - %copy.1550 = bf16[1536,3,128,192]{2,0,3,1:T(8,128)(2,1)} copy(%param_0.1420), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'dense_layers\'][\'self_attention\'][\'wq_b\'][\'kernel\']"} - ROOT %bitcast.673 = bf16[3,1536,128,192]{2,1,3,0:T(8,128)(2,1)} bitcast(%copy.1550), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/dense_layers.wrapped_fn/transpose" stack_frame_id=0} -} - -%region_221.246 (reduce_sum.893: f32[], reduce_sum.603: f32[]) -> f32[] { - %reduce_sum.893 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.603 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.604 = f32[]{:T(128)} add(%reduce_sum.893, %reduce_sum.603), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_187.212 (reduce_sum.655: f32[], reduce_sum.449: f32[]) -> f32[] { - %reduce_sum.655 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.449 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.450 = f32[]{:T(128)} add(%reduce_sum.655, %reduce_sum.449), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.468 (param_0.4136: f32[1536,3,128,192], param_1.5026: f32[], param_2.4295: f32[], param_3.2951: f32[], param_4.2203: f32[1536,3,128,192], param_5.2006: f32[], param_6.1443: f32[3,1536,128,192], param_7.1124: pred[], param_8.889: f32[1536,3,128,192]) -> (f32[], f32[1536,3,128,192], f32[1536,3,128,192], f32[1536,3,128,192], f32[]) { - %param_0.4136 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} parameter(0) - %param_3.2951 = f32[]{:T(128)S(6)} parameter(3) - %mul.4713.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%param_3.2951), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.1124 = pred[]{:T(512)S(6)} parameter(7) - %select_n.2121.clone.1 = pred[1536,3,128,192]{2,3,1,0:T(8,128)(4,1)} broadcast(%param_7.1124), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.1443 = f32[3,1536,128,192]{2,3,0,1:T(8,128)} parameter(6) - %bitcast.1374.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} bitcast(%param_6.1443), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/dense_layers.wrapped_fn/transpose" stack_frame_id=0} - %param_5.2006 = f32[]{:T(128)} parameter(5) - %div.2562.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%param_5.2006), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.2561.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} divide(%bitcast.1374.clone.1, %div.2562.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.2120.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} select(%select_n.2121.clone.1, %bitcast.1374.clone.1, %div.2561.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %constant.4860.clone.1 = f32[]{:T(128)} constant(0.1) - %broadcast.4272.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%constant.4860.clone.1), dimensions={}, metadata={op_name="broadcast.334"} - %mul.4719.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%select_n.2120.clone.1, %broadcast.4272.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.889 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} parameter(8) - %constant.4864.clone.1 = f32[]{:T(128)} constant(0.9) - %mul.4720.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%constant.4864.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.4718.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%param_8.889, %mul.4720.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3488.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} add(%mul.4719.clone.1, %mul.4718.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.4295 = f32[]{:T(128)S(6)} parameter(2) - %div.2558.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%param_2.4295), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %integer_pow.399.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%select_n.2120.clone.1, %select_n.2120.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} - %constant.4863.clone.1 = f32[]{:T(128)} constant(0.05) - %mul.4717.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%constant.4863.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.4715.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%integer_pow.399.clone.1, %mul.4717.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.2203 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} parameter(4) - %constant.4862.clone.1 = f32[]{:T(128)} constant(0.95) - %mul.4716.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%constant.4862.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.4714.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%param_4.2203, %mul.4716.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3487.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} add(%mul.4715.clone.1, %mul.4714.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.5026 = f32[]{:T(128)S(6)} parameter(1) - %div.2557.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%param_1.5026), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.2556.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} divide(%add.3487.clone.1, %div.2557.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %sqrt.157.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} sqrt(%div.2556.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} - %constant.4861.clone.1 = f32[]{:T(128)} constant(1e-08) - %add.3486.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} broadcast(%constant.4861.clone.1), dimensions={}, metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %add.3485.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} add(%sqrt.157.clone.1, %add.3486.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %multiply.1293.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%div.2558.clone.1, %add.3485.clone.1), metadata={op_name="multiply.290"} - %div.2555.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} divide(%add.3488.clone.1, %multiply.1293.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.4712.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%param_0.4136, %broadcast.4272.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3484.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} add(%div.2555.clone.1, %mul.4712.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.4711.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%mul.4713.clone.1, %add.3484.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3483.clone.1 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} add(%param_0.4136, %mul.4711.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.565 = f32[1536,3,128,192]{2,3,1,0:T(8,128)} multiply(%add.3483.clone.1, %add.3483.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.5071 = f32[]{:T(128)} constant(0) - %reduce.670 = f32[]{:T(128)} reduce(%square.565, %constant.5071), dimensions={0,1,2,3}, to_apply=%region_221.246, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.671.clone.1 = f32[]{:T(128)} reduce(%integer_pow.399.clone.1, %constant.5071), dimensions={0,1,2,3}, to_apply=%region_187.212, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.656 = (f32[]{:T(128)}, f32[1536,3,128,192]{2,3,1,0:T(8,128)}, f32[1536,3,128,192]{2,3,1,0:T(8,128)}, f32[1536,3,128,192]{2,3,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.670, %add.3483.clone.1, %add.3487.clone.1, %add.3488.clone.1, %reduce.671.clone.1) -} - -%region_160.185 (reduce_sum.473: f32[], reduce_sum.293: f32[]) -> f32[] { - %reduce_sum.473 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.293 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.300 = f32[]{:T(128)} add(%reduce_sum.473, %reduce_sum.293), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_158.183 (reduce_sum.459: f32[], reduce_sum.460: f32[]) -> f32[] { - %reduce_sum.459 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.460 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.461 = f32[]{:T(128)} add(%reduce_sum.459, %reduce_sum.460), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.494 (param_0.4162: bf16[256,512,512], param_1.5048: bf16[256,512,512]) -> (f32[], f32[]) { - %param_0.4162 = bf16[256,512,512]{2,1,0:T(8,128)(2,1)} parameter(0) - %broadcast_in_dim.1245 = f32[256,512,512]{2,1,0:T(8,128)} convert(%param_0.4162), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/broadcast_in_dim" stack_frame_id=0} - %bitcast.695 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%broadcast_in_dim.1245), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/transpose" stack_frame_id=0} - %square.570 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.695, %bitcast.695), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.5097 = f32[]{:T(128)} constant(0) - %reduce.672 = f32[]{:T(128)} reduce(%square.570, %constant.5097), dimensions={0,1,2,3}, to_apply=%region_160.185, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %param_1.5048 = bf16[256,512,512]{2,1,0:T(8,128)(2,1)} parameter(1) - %broadcast_in_dim.1253.clone.1 = f32[256,512,512]{2,1,0:T(8,128)} convert(%param_1.5048), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/broadcast_in_dim" stack_frame_id=0} - %bitcast.703.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%broadcast_in_dim.1253.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/transpose" stack_frame_id=0} - %square.576.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.703.clone.1, %bitcast.703.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %reduce.674.clone.1 = f32[]{:T(128)} reduce(%square.576.clone.1, %constant.5097), dimensions={0,1,2,3}, to_apply=%region_158.183, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.764 = (f32[]{:T(128)}, f32[]{:T(128)}) tuple(%reduce.672, %reduce.674.clone.1) -} - -%region_159.184 (reduce_sum.466: f32[], reduce_sum.279: f32[]) -> f32[] { - %reduce_sum.466 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.279 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.286 = f32[]{:T(128)} add(%reduce_sum.466, %reduce_sum.279), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.496 (param_0.4161: bf16[256,512,512]) -> f32[] { - %param_0.4161 = bf16[256,512,512]{2,1,0:T(8,128)(2,1)} parameter(0) - %broadcast_in_dim.1249 = f32[256,512,512]{2,1,0:T(8,128)} convert(%param_0.4161), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/broadcast_in_dim" stack_frame_id=0} - %bitcast.699 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%broadcast_in_dim.1249), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/transpose" stack_frame_id=0} - %square.573 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.699, %bitcast.699), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.5096 = f32[]{:T(128)} constant(0) - ROOT %reduce.673 = f32[]{:T(128)} reduce(%square.573, %constant.5096), dimensions={0,1,2,3}, to_apply=%region_159.184, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} -} - -%region_227.252 (reduce_sum.935: f32[], reduce_sum.631: f32[]) -> f32[] { - %reduce_sum.935 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.631 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.632 = f32[]{:T(128)} add(%reduce_sum.935, %reduce_sum.631), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_193.218 (reduce_sum.697: f32[], reduce_sum.471: f32[]) -> f32[] { - %reduce_sum.697 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.471 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.472 = f32[]{:T(128)} add(%reduce_sum.697, %reduce_sum.471), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.514 (param_0.4130: f32[], param_1.5020: f32[256,1,512,512], param_2.4289: f32[], param_3.2945: f32[256,1,512,512], param_4.2197: f32[], param_5.2000: bf16[256,512,512], param_6.1437: pred[], param_7.1118: f32[], param_8.883: f32[256,1,512,512]) -> (f32[], f32[256,1,512,512], f32[256,1,512,512], f32[256,1,512,512], f32[]) { - %param_8.883 = f32[256,1,512,512]{3,2,1,0:T(8,128)} parameter(8) - %bitcast.1359.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%param_8.883), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wo\']"} - %param_7.1118 = f32[]{:T(128)S(6)} parameter(7) - %mul.4662.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_7.1118), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_6.1437 = pred[]{:T(512)S(6)} parameter(6) - %select_n.2103.clone.1 = pred[256,1,512,512]{3,2,0,1:T(8,128)(4,1)} broadcast(%param_6.1437), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_5.2000 = bf16[256,512,512]{2,1,0:T(8,128)(2,1)} parameter(5) - %broadcast_in_dim.1459.clone.1 = f32[256,512,512]{2,1,0:T(8,128)} convert(%param_5.2000), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/broadcast_in_dim" stack_frame_id=0} - %bitcast.1361.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%broadcast_in_dim.1459.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/transpose" stack_frame_id=0} - %param_4.2197 = f32[]{:T(128)} parameter(4) - %div.2520.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_4.2197), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.2519.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} divide(%bitcast.1361.clone.1, %div.2520.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.2102.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} select(%select_n.2103.clone.1, %bitcast.1361.clone.1, %div.2519.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %constant.4830.clone.1 = f32[]{:T(128)} constant(0.1) - %broadcast.4252.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4830.clone.1), dimensions={}, metadata={op_name="broadcast.2362"} - %mul.4664.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%select_n.2102.clone.1, %broadcast.4252.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_3.2945 = f32[256,1,512,512]{3,2,1,0:T(8,128)} parameter(3) - %bitcast.1360.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%param_3.2945), sharding={replicated}, metadata={op_name="state.opt_state[0].mu[\'params\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wo\']"} - %constant.4829.clone.1 = f32[]{:T(128)} constant(0.9) - %broadcast.4251.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4829.clone.1), dimensions={}, metadata={op_name="broadcast.329"} - %mul.4663.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1360.clone.1, %broadcast.4251.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3453.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%mul.4664.clone.1, %mul.4663.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.4289 = f32[]{:T(128)S(6)} parameter(2) - %div.2518.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_2.4289), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %integer_pow.393.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%select_n.2102.clone.1, %select_n.2102.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} - %constant.4828.clone.1 = f32[]{:T(128)} constant(0.05) - %broadcast.4254.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4828.clone.1), dimensions={}, metadata={op_name="broadcast.2365"} - %mul.4666.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%integer_pow.393.clone.1, %broadcast.4254.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_1.5020 = f32[256,1,512,512]{3,2,1,0:T(8,128)} parameter(1) - %bitcast.1362.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%param_1.5020), sharding={replicated}, metadata={op_name="state.opt_state[0].nu[\'params\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wo\']"} - %constant.4827.clone.1 = f32[]{:T(128)} constant(0.95) - %broadcast.4253.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4827.clone.1), dimensions={}, metadata={op_name="broadcast.312"} - %mul.4665.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1362.clone.1, %broadcast.4253.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3454.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%mul.4666.clone.1, %mul.4665.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_0.4130 = f32[]{:T(128)S(6)} parameter(0) - %div.2517.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_0.4130), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.2516.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} divide(%add.3454.clone.1, %div.2517.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %sqrt.151.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} sqrt(%div.2516.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} - %constant.4831.clone.1 = f32[]{:T(128)} constant(1e-08) - %broadcast.4250.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4831.clone.1), dimensions={}, metadata={op_name="broadcast.305"} - %add.3452.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%sqrt.151.clone.1, %broadcast.4250.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %multiply.1287.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%div.2518.clone.1, %add.3452.clone.1), metadata={op_name="multiply.296"} - %div.2515.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} divide(%add.3453.clone.1, %multiply.1287.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.4661.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1359.clone.1, %broadcast.4252.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3451.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%div.2515.clone.1, %mul.4661.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.4660.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%mul.4662.clone.1, %add.3451.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3450.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%bitcast.1359.clone.1, %mul.4660.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.577 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%add.3450.clone.1, %add.3450.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.5065 = f32[]{:T(128)} constant(0) - %reduce.675 = f32[]{:T(128)} reduce(%square.577, %constant.5065), dimensions={0,1,2,3}, to_apply=%region_227.252, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %bitcast.849.clone.1 = f32[256,1,512,512]{3,2,1,0:T(8,128)} bitcast(%add.3454.clone.1) - %bitcast.822.clone.1 = f32[256,1,512,512]{3,2,1,0:T(8,128)} bitcast(%add.3453.clone.1) - %reduce.684.clone.1 = f32[]{:T(128)} reduce(%integer_pow.393.clone.1, %constant.5065), dimensions={0,1,2,3}, to_apply=%region_193.218, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.666 = (f32[]{:T(128)}, f32[256,1,512,512]{3,2,0,1:T(8,128)}, f32[256,1,512,512]{3,2,1,0:T(8,128)}, f32[256,1,512,512]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.675, %add.3450.clone.1, %bitcast.849.clone.1, %bitcast.822.clone.1, %reduce.684.clone.1) -} - -%region_226.251 (reduce_sum.928: f32[], reduce_sum.625: f32[]) -> f32[] { - %reduce_sum.928 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.625 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.626 = f32[]{:T(128)} add(%reduce_sum.928, %reduce_sum.625), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_192.217 (reduce_sum.690: f32[], reduce_sum.465: f32[]) -> f32[] { - %reduce_sum.690 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.465 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.470 = f32[]{:T(128)} add(%reduce_sum.690, %reduce_sum.465), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.515 (param_0.4131: f32[], param_1.5021: f32[256,1,512,512], param_2.4290: f32[], param_3.2946: f32[256,1,512,512], param_4.2198: f32[], param_5.2001: bf16[256,512,512], param_6.1438: pred[], param_7.1119: f32[], param_8.884: f32[256,1,512,512]) -> (f32[], f32[256,1,512,512], f32[256,1,512,512], f32[256,1,512,512], f32[]) { - %param_8.884 = f32[256,1,512,512]{3,2,1,0:T(8,128)} parameter(8) - %bitcast.1363.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%param_8.884), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wi_1\']"} - %param_7.1119 = f32[]{:T(128)S(6)} parameter(7) - %mul.4669.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_7.1119), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_6.1438 = pred[]{:T(512)S(6)} parameter(6) - %select_n.2105.clone.1 = pred[256,1,512,512]{3,2,0,1:T(8,128)(4,1)} broadcast(%param_6.1438), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_5.2001 = bf16[256,512,512]{2,1,0:T(8,128)(2,1)} parameter(5) - %broadcast_in_dim.1460.clone.1 = f32[256,512,512]{2,1,0:T(8,128)} convert(%param_5.2001), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/broadcast_in_dim" stack_frame_id=0} - %bitcast.1365.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%broadcast_in_dim.1460.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/transpose" stack_frame_id=0} - %param_4.2198 = f32[]{:T(128)} parameter(4) - %div.2526.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_4.2198), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.2525.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} divide(%bitcast.1365.clone.1, %div.2526.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.2104.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} select(%select_n.2105.clone.1, %bitcast.1365.clone.1, %div.2525.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %constant.4835.clone.1 = f32[]{:T(128)} constant(0.1) - %broadcast.4257.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4835.clone.1), dimensions={}, metadata={op_name="broadcast.2362"} - %mul.4671.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%select_n.2104.clone.1, %broadcast.4257.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_3.2946 = f32[256,1,512,512]{3,2,1,0:T(8,128)} parameter(3) - %bitcast.1364.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%param_3.2946), sharding={replicated}, metadata={op_name="state.opt_state[0].mu[\'params\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wi_1\']"} - %constant.4834.clone.1 = f32[]{:T(128)} constant(0.9) - %broadcast.4256.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4834.clone.1), dimensions={}, metadata={op_name="broadcast.329"} - %mul.4670.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1364.clone.1, %broadcast.4256.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3458.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%mul.4671.clone.1, %mul.4670.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.4290 = f32[]{:T(128)S(6)} parameter(2) - %div.2524.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_2.4290), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %integer_pow.394.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%select_n.2104.clone.1, %select_n.2104.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} - %constant.4833.clone.1 = f32[]{:T(128)} constant(0.05) - %broadcast.4259.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4833.clone.1), dimensions={}, metadata={op_name="broadcast.2365"} - %mul.4673.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%integer_pow.394.clone.1, %broadcast.4259.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_1.5021 = f32[256,1,512,512]{3,2,1,0:T(8,128)} parameter(1) - %bitcast.1366.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%param_1.5021), sharding={replicated}, metadata={op_name="state.opt_state[0].nu[\'params\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wi_1\']"} - %constant.4832.clone.1 = f32[]{:T(128)} constant(0.95) - %broadcast.4258.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4832.clone.1), dimensions={}, metadata={op_name="broadcast.312"} - %mul.4672.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1366.clone.1, %broadcast.4258.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3459.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%mul.4673.clone.1, %mul.4672.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_0.4131 = f32[]{:T(128)S(6)} parameter(0) - %div.2523.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_0.4131), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.2522.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} divide(%add.3459.clone.1, %div.2523.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %sqrt.152.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} sqrt(%div.2522.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} - %constant.4836.clone.1 = f32[]{:T(128)} constant(1e-08) - %broadcast.4255.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4836.clone.1), dimensions={}, metadata={op_name="broadcast.305"} - %add.3457.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%sqrt.152.clone.1, %broadcast.4255.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %multiply.1288.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%div.2524.clone.1, %add.3457.clone.1), metadata={op_name="multiply.295"} - %div.2521.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} divide(%add.3458.clone.1, %multiply.1288.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.4668.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1363.clone.1, %broadcast.4257.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3456.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%div.2521.clone.1, %mul.4668.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.4667.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%mul.4669.clone.1, %add.3456.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3455.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%bitcast.1363.clone.1, %mul.4667.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.578 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%add.3455.clone.1, %add.3455.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.5066 = f32[]{:T(128)} constant(0) - %reduce.676 = f32[]{:T(128)} reduce(%square.578, %constant.5066), dimensions={0,1,2,3}, to_apply=%region_226.251, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %bitcast.840.clone.1 = f32[256,1,512,512]{3,2,1,0:T(8,128)} bitcast(%add.3459.clone.1) - %bitcast.813.clone.1 = f32[256,1,512,512]{3,2,1,0:T(8,128)} bitcast(%add.3458.clone.1) - %reduce.685.clone.1 = f32[]{:T(128)} reduce(%integer_pow.394.clone.1, %constant.5066), dimensions={0,1,2,3}, to_apply=%region_192.217, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.665 = (f32[]{:T(128)}, f32[256,1,512,512]{3,2,0,1:T(8,128)}, f32[256,1,512,512]{3,2,1,0:T(8,128)}, f32[256,1,512,512]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.676, %add.3455.clone.1, %bitcast.840.clone.1, %bitcast.813.clone.1, %reduce.685.clone.1) -} - -%region_225.250 (reduce_sum.921: f32[], reduce_sum.619: f32[]) -> f32[] { - %reduce_sum.921 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.619 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.624 = f32[]{:T(128)} add(%reduce_sum.921, %reduce_sum.619), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_191.216 (reduce_sum.683: f32[], reduce_sum.463: f32[]) -> f32[] { - %reduce_sum.683 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.463 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.464 = f32[]{:T(128)} add(%reduce_sum.683, %reduce_sum.463), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.516 (param_0.4132: f32[], param_1.5022: f32[256,1,512,512], param_2.4291: f32[], param_3.2947: f32[256,1,512,512], param_4.2199: f32[], param_5.2002: bf16[256,512,512], param_6.1439: pred[], param_7.1120: f32[], param_8.885: f32[256,1,512,512]) -> (f32[], f32[256,1,512,512], f32[256,1,512,512], f32[256,1,512,512], f32[]) { - %param_8.885 = f32[256,1,512,512]{3,2,1,0:T(8,128)} parameter(8) - %bitcast.1367.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%param_8.885), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wi_0\']"} - %param_7.1120 = f32[]{:T(128)S(6)} parameter(7) - %mul.4676.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_7.1120), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_6.1439 = pred[]{:T(512)S(6)} parameter(6) - %select_n.2107.clone.1 = pred[256,1,512,512]{3,2,0,1:T(8,128)(4,1)} broadcast(%param_6.1439), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_5.2002 = bf16[256,512,512]{2,1,0:T(8,128)(2,1)} parameter(5) - %broadcast_in_dim.1461.clone.1 = f32[256,512,512]{2,1,0:T(8,128)} convert(%param_5.2002), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/broadcast_in_dim" stack_frame_id=0} - %bitcast.1369.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%broadcast_in_dim.1461.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/transpose" stack_frame_id=0} - %param_4.2199 = f32[]{:T(128)} parameter(4) - %div.2532.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_4.2199), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.2531.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} divide(%bitcast.1369.clone.1, %div.2532.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.2106.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} select(%select_n.2107.clone.1, %bitcast.1369.clone.1, %div.2531.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %constant.4840.clone.1 = f32[]{:T(128)} constant(0.1) - %broadcast.4262.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4840.clone.1), dimensions={}, metadata={op_name="broadcast.2362"} - %mul.4678.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%select_n.2106.clone.1, %broadcast.4262.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_3.2947 = f32[256,1,512,512]{3,2,1,0:T(8,128)} parameter(3) - %bitcast.1368.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%param_3.2947), sharding={replicated}, metadata={op_name="state.opt_state[0].mu[\'params\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wi_0\']"} - %constant.4839.clone.1 = f32[]{:T(128)} constant(0.9) - %broadcast.4261.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4839.clone.1), dimensions={}, metadata={op_name="broadcast.329"} - %mul.4677.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1368.clone.1, %broadcast.4261.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3463.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%mul.4678.clone.1, %mul.4677.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.4291 = f32[]{:T(128)S(6)} parameter(2) - %div.2530.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_2.4291), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %integer_pow.395.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%select_n.2106.clone.1, %select_n.2106.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} - %constant.4838.clone.1 = f32[]{:T(128)} constant(0.05) - %broadcast.4264.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4838.clone.1), dimensions={}, metadata={op_name="broadcast.2365"} - %mul.4680.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%integer_pow.395.clone.1, %broadcast.4264.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_1.5022 = f32[256,1,512,512]{3,2,1,0:T(8,128)} parameter(1) - %bitcast.1370.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} bitcast(%param_1.5022), sharding={replicated}, metadata={op_name="state.opt_state[0].nu[\'params\'][\'decoder\'][\'moe_layers\'][\'DeepSeekMoeBlock_0\'][\'MoeBlock_0\'][\'wi_0\']"} - %constant.4837.clone.1 = f32[]{:T(128)} constant(0.95) - %broadcast.4263.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4837.clone.1), dimensions={}, metadata={op_name="broadcast.312"} - %mul.4679.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1370.clone.1, %broadcast.4263.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3464.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%mul.4680.clone.1, %mul.4679.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_0.4132 = f32[]{:T(128)S(6)} parameter(0) - %div.2529.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%param_0.4132), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.2528.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} divide(%add.3464.clone.1, %div.2529.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %sqrt.153.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} sqrt(%div.2528.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} - %constant.4841.clone.1 = f32[]{:T(128)} constant(1e-08) - %broadcast.4260.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} broadcast(%constant.4841.clone.1), dimensions={}, metadata={op_name="broadcast.305"} - %add.3462.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%sqrt.153.clone.1, %broadcast.4260.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %multiply.1289.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%div.2530.clone.1, %add.3462.clone.1), metadata={op_name="multiply.294"} - %div.2527.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} divide(%add.3463.clone.1, %multiply.1289.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.4675.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%bitcast.1367.clone.1, %broadcast.4262.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3461.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%div.2527.clone.1, %mul.4675.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.4674.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%mul.4676.clone.1, %add.3461.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3460.clone.1 = f32[256,1,512,512]{3,2,0,1:T(8,128)} add(%bitcast.1367.clone.1, %mul.4674.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.579 = f32[256,1,512,512]{3,2,0,1:T(8,128)} multiply(%add.3460.clone.1, %add.3460.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.5067 = f32[]{:T(128)} constant(0) - %reduce.677 = f32[]{:T(128)} reduce(%square.579, %constant.5067), dimensions={0,1,2,3}, to_apply=%region_225.250, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %bitcast.831.clone.1 = f32[256,1,512,512]{3,2,1,0:T(8,128)} bitcast(%add.3464.clone.1) - %bitcast.804.clone.1 = f32[256,1,512,512]{3,2,1,0:T(8,128)} bitcast(%add.3463.clone.1) - %reduce.686.clone.1 = f32[]{:T(128)} reduce(%integer_pow.395.clone.1, %constant.5067), dimensions={0,1,2,3}, to_apply=%region_191.216, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.664 = (f32[]{:T(128)}, f32[256,1,512,512]{3,2,0,1:T(8,128)}, f32[256,1,512,512]{3,2,1,0:T(8,128)}, f32[256,1,512,512]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.677, %add.3460.clone.1, %bitcast.831.clone.1, %bitcast.804.clone.1, %reduce.686.clone.1) -} - -%region_155.180 (reduce_sum.438: f32[], reduce_sum.259: f32[]) -> f32[] { - %reduce_sum.438 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.259 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.260 = f32[]{:T(128)} add(%reduce_sum.438, %reduce_sum.259), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.528.clone.clone.clone (param_0.4075: bf16[4,128,129280], param_1.4954: s32[4,128], param_2.4222: f32[4,128], param_3.2913: f32[4,128], param_4.2170: bf16[4,128], param_5.1978: f32[4,128]) -> bf16[4,128,129280] { - %param_5.1978 = f32[4,128]{1,0:T(4,128)S(1)} parameter(5) - %mul.4889 = f32[4,128,129280]{2,1,0:T(8,128)} broadcast(%param_5.1978), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - %param_3.2913 = f32[4,128]{1,0:T(4,128)S(1)} parameter(3) - %mul.4888 = f32[4,128,129280]{2,1,0:T(8,128)} broadcast(%param_3.2913), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - %param_0.4075 = bf16[4,128,129280]{2,1,0:T(8,128)(2,1)} parameter(0) - %convert_element_type.3151 = f32[4,128,129280]{2,1,0:T(8,128)} convert(%param_0.4075), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %param_4.2170 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(4) - %sub.791 = f32[4,128,129280]{2,1,0:T(8,128)} broadcast(%param_4.2170), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %sub.790 = f32[4,128,129280]{2,1,0:T(8,128)} subtract(%convert_element_type.3151, %sub.791), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %exp.534 = f32[4,128,129280]{2,1,0:T(8,128)} exponential(%sub.790), metadata={op_name="jit(train_step)/jvp()/exp" stack_frame_id=0} - %mul.4887 = f32[4,128,129280]{2,1,0:T(8,128)} multiply(%mul.4888, %exp.534), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - %param_2.4222 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %div.2685 = f32[4,128,129280]{2,1,0:T(8,128)} broadcast(%param_2.4222), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} - %div.2684 = f32[4,128,129280]{2,1,0:T(8,128)} divide(%mul.4887, %div.2685), metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} - %param_1.4954 = s32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %eq.363 = s32[4,128,129280]{2,1,0:T(8,128)} broadcast(%param_1.4954), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} - %eq.362 = s32[4,128,129280]{2,1,0:T(8,128)} iota(), iota_dimension=2, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} - %eq.361 = pred[4,128,129280]{2,1,0:T(8,128)(4,1)} compare(%eq.363, %eq.362), direction=EQ, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} - %convert_element_type.3150 = f32[4,128,129280]{2,1,0:T(8,128)} convert(%eq.361), metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/convert_element_type" stack_frame_id=0} - %sub.789 = f32[4,128,129280]{2,1,0:T(8,128)} subtract(%div.2684, %convert_element_type.3150), metadata={op_name="jit(train_step)/transpose(jvp())/sub" stack_frame_id=0} - %mul.4886 = f32[4,128,129280]{2,1,0:T(8,128)} multiply(%mul.4889, %sub.789), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - ROOT %convert_element_type.3149 = bf16[4,128,129280]{2,1,0:T(8,128)(2,1)} convert(%mul.4886), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} -} - -%fused_computation.935.clone.clone (param_0.4076: f32[4,128], param_1.4955: bf16[4,128,512], param_2.4224: bf16[512]) -> bf16[4,128,512] { - %param_2.4224 = bf16[512]{0:T(512)(128)(2,1)S(1)} parameter(2) - %dot_general.831 = bf16[4,128,512]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.4224), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.4955 = bf16[4,128,512]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %convert_element_type.3153 = f32[4,128,512]{2,1,0:T(8,128)} convert(%param_1.4955), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %param_0.4076 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.4891 = f32[4,128,512]{2,1,0:T(8,128)} broadcast(%param_0.4076), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %mul.4890 = f32[4,128,512]{2,1,0:T(8,128)} multiply(%convert_element_type.3153, %mul.4891), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %convert_element_type.3152 = bf16[4,128,512]{2,1,0:T(8,128)(2,1)} convert(%mul.4890), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - ROOT %dot_general.830 = bf16[4,128,512]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.831, %convert_element_type.3152), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} -} - -%fused_computation.517 (param_0.4165: bf16[4,128,129280], param_1.5050: s32[4,128], param_2.4316: f32[4,128], param_3.2969: f32[4,128], param_4.2219: bf16[4,128], param_5.2020: f32[4,128], param_6.1457: f32[4,128], param_7.1138: bf16[4,128,512], param_8.902: bf16[512]) -> (f32[], bf16[512,129280,1]) { - %param_6.1457 = f32[4,128]{1,0:T(4,128)S(1)} parameter(6) - %param_7.1138 = bf16[4,128,512]{2,1,0:T(8,128)(2,1)S(1)} parameter(7) - %param_8.902 = bf16[512]{0:T(512)(128)(2,1)S(1)} parameter(8) - %fusion.573.clone.1 = bf16[4,128,512]{2,1,0:T(8,128)(2,1)} fusion(%param_6.1457, %param_7.1138, %param_8.902), kind=kLoop, calls=%fused_computation.935.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %param_0.4165 = bf16[4,128,129280]{2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.5050 = s32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %param_2.4316 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %param_3.2969 = f32[4,128]{1,0:T(4,128)S(1)} parameter(3) - %param_4.2219 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(4) - %param_5.2020 = f32[4,128]{1,0:T(4,128)S(1)} parameter(5) - %multiply_convert_fusion.1.clone.1 = bf16[4,128,129280]{2,1,0:T(8,128)(2,1)} fusion(%param_0.4165, %param_1.5050, %param_2.4316, %param_3.2969, %param_4.2219, /*index=5*/%param_5.2020), kind=kLoop, calls=%fused_computation.528.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %convolution.141.clone.1 = bf16[512,129280,1]{1,0,2:T(8,128)(2,1)} convolution(%fusion.573.clone.1, %multiply_convert_fusion.1.clone.1), window={size=4}, dim_labels=0fb_0io->bf0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/dot_general" stack_frame_id=0} - %bitcast.776 = bf16[512,129280]{1,0:T(8,128)(2,1)} bitcast(%convolution.141.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/dot_general" stack_frame_id=0} - %convert_element_type.2653 = f32[512,129280]{1,0:T(8,128)} convert(%bitcast.776), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} - %square.581 = f32[512,129280]{1,0:T(8,128)} multiply(%convert_element_type.2653, %convert_element_type.2653), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.5100 = f32[]{:T(128)} constant(0) - %reduce.678 = f32[]{:T(128)} reduce(%square.581, %constant.5100), dimensions={0,1}, to_apply=%region_155.180, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.754 = (f32[]{:T(128)}, bf16[512,129280,1]{1,0,2:T(8,128)(2,1)}) tuple(%reduce.678, %convolution.141.clone.1) -} - -%region_174.199 (reduce_sum.564: f32[], reduce_sum.387: f32[]) -> f32[] { - %reduce_sum.564 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.387 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.388 = f32[]{:T(128)} add(%reduce_sum.564, %reduce_sum.387), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.518 (param_0.4149: bf16[129280,512]) -> f32[] { - %param_0.4149 = bf16[129280,512]{1,0:T(8,128)(2,1)} parameter(0) - %convert_element_type.2655 = f32[129280,512]{1,0:T(8,128)} convert(%param_0.4149), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} - %square.583 = f32[129280,512]{1,0:T(8,128)} multiply(%convert_element_type.2655, %convert_element_type.2655), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.5084 = f32[]{:T(128)} constant(0) - ROOT %reduce.679 = f32[]{:T(128)} reduce(%square.583, %constant.5084), dimensions={0,1}, to_apply=%region_174.199, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} -} - -%region_240.265 (reduce_sum.1026: f32[], reduce_sum.689: f32[]) -> f32[] { - %reduce_sum.1026 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.689 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.694 = f32[]{:T(128)} add(%reduce_sum.1026, %reduce_sum.689), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_206.231 (reduce_sum.788: f32[], reduce_sum.533: f32[]) -> f32[] { - %reduce_sum.788 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.533 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.534 = f32[]{:T(128)} add(%reduce_sum.788, %reduce_sum.533), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.519 (param_0.4117: f32[129280,512], param_1.5007: f32[], param_2.4276: f32[], param_3.2932: f32[], param_4.2184: f32[129280,512], param_5.1987: f32[], param_6.1424: bf16[129280,512], param_7.1105: pred[], param_8.870: f32[129280,512]) -> (f32[], f32[129280,512], f32[129280,512], f32[129280,512], f32[]) { - %param_0.4117 = f32[129280,512]{1,0:T(8,128)} parameter(0) - %param_3.2932 = f32[]{:T(128)S(6)} parameter(3) - %mul.4550.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%param_3.2932), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.1105 = pred[]{:T(512)S(6)} parameter(7) - %select_n.2061.clone.1 = pred[129280,512]{1,0:T(8,128)(4,1)} broadcast(%param_7.1105), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.1424 = bf16[129280,512]{1,0:T(8,128)(2,1)} parameter(6) - %convert_element_type.3094.clone.1 = f32[129280,512]{1,0:T(8,128)} convert(%param_6.1424), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} - %param_5.1987 = f32[]{:T(128)} parameter(5) - %div.2426.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%param_5.1987), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.2425.clone.1 = f32[129280,512]{1,0:T(8,128)} divide(%convert_element_type.3094.clone.1, %div.2426.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.2060.clone.1 = f32[129280,512]{1,0:T(8,128)} select(%select_n.2061.clone.1, %convert_element_type.3094.clone.1, %div.2425.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %constant.4750.clone.1 = f32[]{:T(128)} constant(0.1) - %broadcast.4202.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%constant.4750.clone.1), dimensions={}, metadata={op_name="broadcast.318"} - %mul.4556.clone.1 = f32[129280,512]{1,0:T(8,128)} multiply(%select_n.2060.clone.1, %broadcast.4202.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.870 = f32[129280,512]{1,0:T(8,128)} parameter(8) - %constant.4754.clone.1 = f32[]{:T(128)} constant(0.9) - %mul.4557.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%constant.4754.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.4555.clone.1 = f32[129280,512]{1,0:T(8,128)} multiply(%param_8.870, %mul.4557.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3383.clone.1 = f32[129280,512]{1,0:T(8,128)} add(%mul.4556.clone.1, %mul.4555.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.4276 = f32[]{:T(128)S(6)} parameter(2) - %div.2422.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%param_2.4276), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %integer_pow.380.clone.1 = f32[129280,512]{1,0:T(8,128)} multiply(%select_n.2060.clone.1, %select_n.2060.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} - %constant.4753.clone.1 = f32[]{:T(128)} constant(0.05) - %mul.4554.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%constant.4753.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.4552.clone.1 = f32[129280,512]{1,0:T(8,128)} multiply(%integer_pow.380.clone.1, %mul.4554.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.2184 = f32[129280,512]{1,0:T(8,128)} parameter(4) - %constant.4752.clone.1 = f32[]{:T(128)} constant(0.95) - %mul.4553.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%constant.4752.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.4551.clone.1 = f32[129280,512]{1,0:T(8,128)} multiply(%param_4.2184, %mul.4553.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3382.clone.1 = f32[129280,512]{1,0:T(8,128)} add(%mul.4552.clone.1, %mul.4551.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.5007 = f32[]{:T(128)S(6)} parameter(1) - %div.2421.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%param_1.5007), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.2420.clone.1 = f32[129280,512]{1,0:T(8,128)} divide(%add.3382.clone.1, %div.2421.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %sqrt.138.clone.1 = f32[129280,512]{1,0:T(8,128)} sqrt(%div.2420.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} - %constant.4751.clone.1 = f32[]{:T(128)} constant(1e-08) - %add.3381.clone.1 = f32[129280,512]{1,0:T(8,128)} broadcast(%constant.4751.clone.1), dimensions={}, metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %add.3380.clone.1 = f32[129280,512]{1,0:T(8,128)} add(%sqrt.138.clone.1, %add.3381.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %multiply.1274.clone.1 = f32[129280,512]{1,0:T(8,128)} multiply(%div.2422.clone.1, %add.3380.clone.1), metadata={op_name="multiply.309"} - %div.2419.clone.1 = f32[129280,512]{1,0:T(8,128)} divide(%add.3383.clone.1, %multiply.1274.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.4549.clone.1 = f32[129280,512]{1,0:T(8,128)} multiply(%param_0.4117, %broadcast.4202.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3379.clone.1 = f32[129280,512]{1,0:T(8,128)} add(%div.2419.clone.1, %mul.4549.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.4548.clone.1 = f32[129280,512]{1,0:T(8,128)} multiply(%mul.4550.clone.1, %add.3379.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3378.clone.1 = f32[129280,512]{1,0:T(8,128)} add(%param_0.4117, %mul.4548.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.584 = f32[129280,512]{1,0:T(8,128)} multiply(%add.3378.clone.1, %add.3378.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.5052 = f32[]{:T(128)} constant(0) - %reduce.680 = f32[]{:T(128)} reduce(%square.584, %constant.5052), dimensions={0,1}, to_apply=%region_240.265, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.687.clone.1 = f32[]{:T(128)} reduce(%integer_pow.380.clone.1, %constant.5052), dimensions={0,1}, to_apply=%region_206.231, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.667 = (f32[]{:T(128)}, f32[129280,512]{1,0:T(8,128)}, f32[129280,512]{1,0:T(8,128)}, f32[129280,512]{1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.680, %add.3378.clone.1, %add.3382.clone.1, %add.3383.clone.1, %reduce.687.clone.1) -} - -%region_222.247 (reduce_sum.900: f32[], reduce_sum.605: f32[]) -> f32[] { - %reduce_sum.900 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.605 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.610 = f32[]{:T(128)} add(%reduce_sum.900, %reduce_sum.605), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_188.213 (reduce_sum.662: f32[], reduce_sum.451: f32[]) -> f32[] { - %reduce_sum.662 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.451 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.455 = f32[]{:T(128)} add(%reduce_sum.662, %reduce_sum.451), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.520 (param_0.4135: f32[512,129280], param_1.5025: f32[], param_2.4294: f32[], param_3.2950: f32[], param_4.2202: f32[512,129280], param_5.2005: f32[], param_6.1442: bf16[512,129280,1], param_7.1123: pred[], param_8.888: f32[512,129280]) -> (f32[], f32[512,129280], f32[512,129280], f32[512,129280], f32[]) { - %param_0.4135 = f32[512,129280]{1,0:T(8,128)} parameter(0) - %param_3.2950 = f32[]{:T(128)S(6)} parameter(3) - %mul.4703.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%param_3.2950), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.1123 = pred[]{:T(512)S(6)} parameter(7) - %select_n.2117.clone.1 = pred[512,129280]{1,0:T(8,128)(4,1)} broadcast(%param_7.1123), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.1442 = bf16[512,129280,1]{1,0,2:T(8,128)(2,1)} parameter(6) - %bitcast.1372.clone.1 = bf16[512,129280]{1,0:T(8,128)(2,1)} bitcast(%param_6.1442), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/dot_general" stack_frame_id=0} - %convert_element_type.3096.clone.1 = f32[512,129280]{1,0:T(8,128)} convert(%bitcast.1372.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} - %param_5.2005 = f32[]{:T(128)} parameter(5) - %div.2554.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%param_5.2005), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.2553.clone.1 = f32[512,129280]{1,0:T(8,128)} divide(%convert_element_type.3096.clone.1, %div.2554.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.2116.clone.1 = f32[512,129280]{1,0:T(8,128)} select(%select_n.2117.clone.1, %convert_element_type.3096.clone.1, %div.2553.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %constant.4854.clone.1 = f32[]{:T(128)} constant(0.1) - %broadcast.4270.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%constant.4854.clone.1), dimensions={}, metadata={op_name="broadcast.333"} - %mul.4709.clone.1 = f32[512,129280]{1,0:T(8,128)} multiply(%select_n.2116.clone.1, %broadcast.4270.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.888 = f32[512,129280]{1,0:T(8,128)} parameter(8) - %constant.4858.clone.1 = f32[]{:T(128)} constant(0.9) - %mul.4710.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%constant.4858.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.4708.clone.1 = f32[512,129280]{1,0:T(8,128)} multiply(%param_8.888, %mul.4710.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3482.clone.1 = f32[512,129280]{1,0:T(8,128)} add(%mul.4709.clone.1, %mul.4708.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.4294 = f32[]{:T(128)S(6)} parameter(2) - %div.2550.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%param_2.4294), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %integer_pow.398.clone.1 = f32[512,129280]{1,0:T(8,128)} multiply(%select_n.2116.clone.1, %select_n.2116.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} - %constant.4857.clone.1 = f32[]{:T(128)} constant(0.05) - %mul.4707.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%constant.4857.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.4705.clone.1 = f32[512,129280]{1,0:T(8,128)} multiply(%integer_pow.398.clone.1, %mul.4707.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.2202 = f32[512,129280]{1,0:T(8,128)} parameter(4) - %constant.4856.clone.1 = f32[]{:T(128)} constant(0.95) - %mul.4706.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%constant.4856.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.4704.clone.1 = f32[512,129280]{1,0:T(8,128)} multiply(%param_4.2202, %mul.4706.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3481.clone.1 = f32[512,129280]{1,0:T(8,128)} add(%mul.4705.clone.1, %mul.4704.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.5025 = f32[]{:T(128)S(6)} parameter(1) - %div.2549.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%param_1.5025), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.2548.clone.1 = f32[512,129280]{1,0:T(8,128)} divide(%add.3481.clone.1, %div.2549.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %sqrt.156.clone.1 = f32[512,129280]{1,0:T(8,128)} sqrt(%div.2548.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} - %constant.4855.clone.1 = f32[]{:T(128)} constant(1e-08) - %add.3480.clone.1 = f32[512,129280]{1,0:T(8,128)} broadcast(%constant.4855.clone.1), dimensions={}, metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %add.3479.clone.1 = f32[512,129280]{1,0:T(8,128)} add(%sqrt.156.clone.1, %add.3480.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %multiply.1292.clone.1 = f32[512,129280]{1,0:T(8,128)} multiply(%div.2550.clone.1, %add.3479.clone.1), metadata={op_name="multiply.291"} - %div.2547.clone.1 = f32[512,129280]{1,0:T(8,128)} divide(%add.3482.clone.1, %multiply.1292.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.4702.clone.1 = f32[512,129280]{1,0:T(8,128)} multiply(%param_0.4135, %broadcast.4270.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3478.clone.1 = f32[512,129280]{1,0:T(8,128)} add(%div.2547.clone.1, %mul.4702.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.4701.clone.1 = f32[512,129280]{1,0:T(8,128)} multiply(%mul.4703.clone.1, %add.3478.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3477.clone.1 = f32[512,129280]{1,0:T(8,128)} add(%param_0.4135, %mul.4701.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.585 = f32[512,129280]{1,0:T(8,128)} multiply(%add.3477.clone.1, %add.3477.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.5070 = f32[]{:T(128)} constant(0) - %reduce.681 = f32[]{:T(128)} reduce(%square.585, %constant.5070), dimensions={0,1}, to_apply=%region_222.247, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.688.clone.1 = f32[]{:T(128)} reduce(%integer_pow.398.clone.1, %constant.5070), dimensions={0,1}, to_apply=%region_188.213, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.668 = (f32[]{:T(128)}, f32[512,129280]{1,0:T(8,128)}, f32[512,129280]{1,0:T(8,128)}, f32[512,129280]{1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.681, %add.3477.clone.1, %add.3481.clone.1, %add.3482.clone.1, %reduce.688.clone.1) -} - -%region_207.232 (reduce_sum.795: f32[], reduce_sum.535: f32[]) -> f32[] { - %reduce_sum.795 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} - %reduce_sum.535 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} - ROOT %reduce_sum.540 = f32[]{:T(128)} add(%reduce_sum.795, %reduce_sum.535), metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.521 (param_0.4186: bf16[4,128,129280], param_1.5064: f32[4,128], param_2.4326: s32[4,128], param_3.2977: bf16[4,128]) -> f32[4,128] { - %param_2.4326 = s32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %eq.299 = s32[4,128,129280]{2,1,0:T(8,128)} broadcast(%param_2.4326), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} - %eq.294 = s32[4,128,129280]{2,1,0:T(8,128)} iota(), iota_dimension=2, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} - %eq.293 = pred[4,128,129280]{2,1,0:T(8,128)(4,1)} compare(%eq.299, %eq.294), direction=EQ, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} - %param_0.4186 = bf16[4,128,129280]{2,1,0:T(8,128)(2,1)} parameter(0) - %convert_element_type.2660 = f32[4,128,129280]{2,1,0:T(8,128)} convert(%param_0.4186), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %param_3.2977 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(3) - %sub.652 = f32[4,128,129280]{2,1,0:T(8,128)} broadcast(%param_3.2977), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %sub.643 = f32[4,128,129280]{2,1,0:T(8,128)} subtract(%convert_element_type.2660, %sub.652), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %param_1.5064 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %sub.650 = f32[4,128,129280]{2,1,0:T(8,128)} broadcast(%param_1.5064), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %sub.639 = f32[4,128,129280]{2,1,0:T(8,128)} subtract(%sub.643, %sub.650), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %constant.5124 = f32[]{:T(128)} constant(0) - %broadcast.3777 = f32[4,128,129280]{2,1,0:T(8,128)} broadcast(%constant.5124), dimensions={}, metadata={op_name="broadcast.514"} - %mul.3612 = f32[4,128,129280]{2,1,0:T(8,128)} select(%eq.293, %sub.639, %broadcast.3777), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} - ROOT %reduce.682 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%mul.3612, %constant.5124), dimensions={2}, to_apply=%region_207.232, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} -} - -%region_37.47 (reduce_sum.76: f32[], reduce_sum.80: f32[]) -> f32[] { - %reduce_sum.76 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} - %reduce_sum.80 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} - ROOT %reduce_sum.83 = f32[]{:T(128)} add(%reduce_sum.76, %reduce_sum.80), metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.532 (param_0.4187: bf16[4,128,129280], param_1.5065: bf16[4,128]) -> f32[4,128] { - %param_0.4187 = bf16[4,128,129280]{2,1,0:T(8,128)(2,1)} parameter(0) - %convert_element_type.2666 = f32[4,128,129280]{2,1,0:T(8,128)} convert(%param_0.4187), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %param_1.5065 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(1) - %sub.653 = f32[4,128,129280]{2,1,0:T(8,128)} broadcast(%param_1.5065), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %sub.649 = f32[4,128,129280]{2,1,0:T(8,128)} subtract(%convert_element_type.2666, %sub.653), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %exp.448 = f32[4,128,129280]{2,1,0:T(8,128)} exponential(%sub.649), metadata={op_name="jit(train_step)/jvp()/exp" stack_frame_id=0} - %constant.5125 = f32[]{:T(128)} constant(0) - ROOT %reduce.683 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%exp.448, %constant.5125), dimensions={2}, to_apply=%region_37.47, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} -} - -%region_152.177 (reduce_sum.417: f32[], reduce_sum.244: f32[]) -> f32[] { - %reduce_sum.417 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.244 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.251 = f32[]{:T(128)} add(%reduce_sum.417, %reduce_sum.244), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.540 (param_0.4168: f32[3,512,128,256]) -> f32[] { - %param_0.4168 = f32[3,512,128,256]{3,2,0,1:T(8,128)} parameter(0) - %bitcast.752 = f32[512,3,128,256]{3,2,1,0:T(8,128)} bitcast(%param_0.4168), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/dense_layers.wrapped_fn/transpose" stack_frame_id=0} - %square.588 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%bitcast.752, %bitcast.752), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.5103 = f32[]{:T(128)} constant(0) - ROOT %reduce.689 = f32[]{:T(128)} reduce(%square.588, %constant.5103), dimensions={0,1,2,3}, to_apply=%region_152.177, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} -} - -%fused_computation.541 (param_0.1601: f32[512,3,128,256]) -> bf16[3,512,128,256] { - %param_0.1601 = f32[512,3,128,256]{3,2,1,0:T(8,128)} parameter(0) - %copy.1551 = bf16[512,3,128,256]{3,0,2,1:T(8,128)(2,1)} copy(%param_0.1601), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'dense_layers\'][\'self_attention\'][\'wkv_b\'][\'kernel\']"} - ROOT %bitcast.753 = bf16[3,512,128,256]{3,1,2,0:T(8,128)(2,1)} bitcast(%copy.1551), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/dense_layers.wrapped_fn/transpose" stack_frame_id=0} -} - -%region_219.244 (reduce_sum.879: f32[], reduce_sum.591: f32[]) -> f32[] { - %reduce_sum.879 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.591 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.596 = f32[]{:T(128)} add(%reduce_sum.879, %reduce_sum.591), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_185.210 (reduce_sum.641: f32[], reduce_sum.437: f32[]) -> f32[] { - %reduce_sum.641 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.437 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.442 = f32[]{:T(128)} add(%reduce_sum.641, %reduce_sum.437), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.542 (param_0.4138: f32[512,3,128,256], param_1.5028: f32[], param_2.4297: f32[], param_3.2953: f32[], param_4.2205: f32[512,3,128,256], param_5.2008: f32[], param_6.1445: f32[3,512,128,256], param_7.1126: pred[], param_8.891: f32[512,3,128,256]) -> (f32[], f32[512,3,128,256], f32[512,3,128,256], f32[512,3,128,256], f32[]) { - %param_0.4138 = f32[512,3,128,256]{3,2,1,0:T(8,128)} parameter(0) - %param_3.2953 = f32[]{:T(128)S(6)} parameter(3) - %mul.4733.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%param_3.2953), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.1126 = pred[]{:T(512)S(6)} parameter(7) - %select_n.2129.clone.1 = pred[512,3,128,256]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.1126), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.1445 = f32[3,512,128,256]{3,2,0,1:T(8,128)} parameter(6) - %bitcast.1378.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} bitcast(%param_6.1445), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/dense_layers.wrapped_fn/transpose" stack_frame_id=0} - %param_5.2008 = f32[]{:T(128)} parameter(5) - %div.2578.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%param_5.2008), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.2577.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} divide(%bitcast.1378.clone.1, %div.2578.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.2128.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} select(%select_n.2129.clone.1, %bitcast.1378.clone.1, %div.2577.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %constant.4872.clone.1 = f32[]{:T(128)} constant(0.1) - %broadcast.4276.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%constant.4872.clone.1), dimensions={}, metadata={op_name="broadcast.336"} - %mul.4739.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%select_n.2128.clone.1, %broadcast.4276.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.891 = f32[512,3,128,256]{3,2,1,0:T(8,128)} parameter(8) - %constant.4876.clone.1 = f32[]{:T(128)} constant(0.9) - %mul.4740.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%constant.4876.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.4738.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%param_8.891, %mul.4740.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3500.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} add(%mul.4739.clone.1, %mul.4738.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.4297 = f32[]{:T(128)S(6)} parameter(2) - %div.2574.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%param_2.4297), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %integer_pow.401.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%select_n.2128.clone.1, %select_n.2128.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} - %constant.4875.clone.1 = f32[]{:T(128)} constant(0.05) - %mul.4737.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%constant.4875.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.4735.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%integer_pow.401.clone.1, %mul.4737.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.2205 = f32[512,3,128,256]{3,2,1,0:T(8,128)} parameter(4) - %constant.4874.clone.1 = f32[]{:T(128)} constant(0.95) - %mul.4736.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%constant.4874.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.4734.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%param_4.2205, %mul.4736.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3499.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} add(%mul.4735.clone.1, %mul.4734.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.5028 = f32[]{:T(128)S(6)} parameter(1) - %div.2573.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%param_1.5028), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.2572.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} divide(%add.3499.clone.1, %div.2573.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %sqrt.159.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} sqrt(%div.2572.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} - %constant.4873.clone.1 = f32[]{:T(128)} constant(1e-08) - %add.3498.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} broadcast(%constant.4873.clone.1), dimensions={}, metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %add.3497.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} add(%sqrt.159.clone.1, %add.3498.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %multiply.1295.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%div.2574.clone.1, %add.3497.clone.1), metadata={op_name="multiply.288"} - %div.2571.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} divide(%add.3500.clone.1, %multiply.1295.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.4732.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%param_0.4138, %broadcast.4276.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3496.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} add(%div.2571.clone.1, %mul.4732.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.4731.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%mul.4733.clone.1, %add.3496.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.3495.clone.1 = f32[512,3,128,256]{3,2,1,0:T(8,128)} add(%param_0.4138, %mul.4731.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.589 = f32[512,3,128,256]{3,2,1,0:T(8,128)} multiply(%add.3495.clone.1, %add.3495.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.5073 = f32[]{:T(128)} constant(0) - %reduce.690 = f32[]{:T(128)} reduce(%square.589, %constant.5073), dimensions={0,1,2,3}, to_apply=%region_219.244, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.691.clone.1 = f32[]{:T(128)} reduce(%integer_pow.401.clone.1, %constant.5073), dimensions={0,1,2,3}, to_apply=%region_185.210, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.663 = (f32[]{:T(128)}, f32[512,3,128,256]{3,2,1,0:T(8,128)}, f32[512,3,128,256]{3,2,1,0:T(8,128)}, f32[512,3,128,256]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.690, %add.3495.clone.1, %add.3499.clone.1, %add.3500.clone.1, %reduce.691.clone.1) -} - -%region_172.197 (reduce_sum.557: f32[], reduce_sum.381: f32[]) -> f32[] { - %reduce_sum.557 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.381 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.386 = f32[]{:T(128)} add(%reduce_sum.557, %reduce_sum.381), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.782.clone.clone (param_0.4102: f32[4,128], param_1.4999: bf16[4,128,1536], param_2.4258: bf16[1536]) -> bf16[4,128,1536,1] { - %param_2.4258 = bf16[1536]{0:T(1024)(128)(2,1)S(1)} parameter(2) - %dot_general.851 = bf16[4,128,1536]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.4258), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.4999 = bf16[4,128,1536]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %convert_element_type.3175 = f32[4,128,1536]{2,1,0:T(8,128)} convert(%param_1.4999), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/convert_element_type" stack_frame_id=0} - %param_0.4102 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.4937 = f32[4,128,1536]{2,1,0:T(8,128)} broadcast(%param_0.4102), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/mul" stack_frame_id=0} - %mul.4936 = f32[4,128,1536]{2,1,0:T(8,128)} multiply(%convert_element_type.3175, %mul.4937), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/mul" stack_frame_id=0} - %convert_element_type.3174 = bf16[4,128,1536]{2,1,0:T(8,128)(2,1)} convert(%mul.4936), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/convert_element_type" stack_frame_id=0} - %dot_general.850 = bf16[4,128,1536]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.851, %convert_element_type.3174), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/...k,k->...k/dot_general" stack_frame_id=0} - ROOT %bitcast.1466 = bf16[4,128,1536,1]{2,1,0,3:T(8,128)(2,1)} bitcast(%dot_general.850), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/...k,k->...k/dot_general" stack_frame_id=0} -} - -%bitcast_fusion.12 (bitcast_input.12: bf16[4,128,128,192]) -> bf16[4,128,128,192] { - %bitcast_input.12 = bf16[4,128,128,192]{2,1,0,3:T(8,128)(2,1)S(1)} parameter(0) - ROOT %bitcast.1488 = bf16[4,128,128,192]{2,1,0,3:T(8,128)(2,1)} bitcast(%bitcast_input.12) -} - -%fused_computation.551 (param_0.4150: bf16[4,128,128,192], param_1.5039: f32[4,128], param_2.4308: bf16[4,128,1536], param_3.2964: bf16[1536]) -> (f32[], bf16[1536,128,192,1]) { - %param_1.5039 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %param_2.4308 = bf16[4,128,1536]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) - %param_3.2964 = bf16[1536]{0:T(1024)(128)(2,1)S(1)} parameter(3) - %fusion.459.clone.1 = bf16[4,128,1536,1]{2,1,0,3:T(8,128)(2,1)} fusion(%param_1.5039, %param_2.4308, %param_3.2964), kind=kLoop, calls=%fused_computation.782.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/moe_layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_0.4150 = bf16[4,128,128,192]{2,1,0,3:T(8,128)(2,1)S(1)} parameter(0) - %fusion.746 = bf16[4,128,128,192]{2,1,0,3:T(8,128)(2,1)} fusion(%param_0.4150), kind=kLoop, calls=%bitcast_fusion.12 - %convolution.146.clone.1 = bf16[1536,128,192,1]{1,0,3,2:T(8,128)(2,1)} convolution(%fusion.459.clone.1, %fusion.746), window={size=192x4 pad=191_191x0_0 rhs_reversal=1x0}, dim_labels=1fb0_1io0->bf01, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/dot_general" stack_frame_id=0} - %bitcast.861 = bf16[1536,128,192]{1,0,2:T(8,128)(2,1)} bitcast(%convolution.146.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/closed_call/checkpoint/moe_layers/dot_general" stack_frame_id=0} - %broadcast_in_dim.1275 = f32[1536,128,192]{1,0,2:T(8,128)} convert(%bitcast.861), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/while/body/broadcast_in_dim" stack_frame_id=0} - %bitcast.763 = f32[1536,1,128,192]{2,0,3,1:T(8,128)} bitcast(%broadcast_in_dim.1275), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/moe_layers.wrapped_fn/transpose" stack_frame_id=0} - %square.592 = f32[1536,1,128,192]{2,0,3,1:T(8,128)} multiply(%bitcast.763, %bitcast.763), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.5085 = f32[]{:T(128)} constant(0) - %reduce.692 = f32[]{:T(128)} reduce(%square.592, %constant.5085), dimensions={0,1,2,3}, to_apply=%region_172.197, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.763 = (f32[]{:T(128)}, bf16[1536,128,192,1]{1,0,3,2:T(8,128)(2,1)}) tuple(%reduce.692, %convolution.146.clone.1) -} - -%region_239.264 (reduce_sum.1019: f32[], reduce_sum.687: f32[]) -> f32[] { - %reduce_sum.1019 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.687 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.688 = f32[]{:T(128)} add(%reduce_sum.1019, %reduce_sum.687), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_205.230 (reduce_sum.781: f32[], reduce_sum.527: f32[]) -> f32[] { - %reduce_sum.781 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.527 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.528 = f32[]{:T(128)} add(%reduce_sum.781, %reduce_sum.527), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.556 (param_0.4118: f32[], param_1.5008: f32[], param_2.4277: f32[], param_3.2933: f32[1536,1,128,192], param_4.2185: f32[1536,1,128,192], param_5.1988: f32[], param_6.1425: bf16[1536,128,192,1], param_7.1106: pred[], param_8.871: f32[1536,1,128,192]) -> (f32[], f32[1536,1,128,192], f32[1536,1,128,192], f32[1536,1,128,192], f32[]) { diff --git a/tests/utils/reference_hlo_llama3_8b.txt b/tests/utils/reference_hlo_llama3_8b.txt index 27c6529df2..e9609981a2 100644 --- a/tests/utils/reference_hlo_llama3_8b.txt +++ b/tests/utils/reference_hlo_llama3_8b.txt @@ -1,4 +1,4 @@ -HloModule jit_train_step, is_scheduled=true, input_output_alias={ {0}: (0, {}, may-alias), {1}: (1, {}, may-alias), {2}: (2, {}, may-alias), {3}: (3, {}, may-alias), {4}: (4, {}, may-alias), {5}: (5, {}, may-alias), {6}: (6, {}, may-alias), {7}: (7, {}, may-alias), {8}: (8, {}, may-alias), {9}: (9, {}, may-alias), {10}: (10, {}, may-alias), {11}: (11, {}, may-alias), {12}: (12, {}, may-alias), {13}: (13, {}, may-alias), {14}: (14, {}, may-alias), {15}: (15, {}, may-alias), {16}: (16, {}, may-alias), {17}: (17, {}, may-alias), {18}: (18, {}, may-alias), {19}: (19, {}, may-alias), {20}: (20, {}, may-alias), {21}: (21, {}, may-alias), {22}: (22, {}, may-alias), {23}: (23, {}, may-alias), {24}: (24, {}, may-alias), {25}: (25, {}, may-alias), {26}: (26, {}, may-alias), {27}: (27, {}, may-alias), {28}: (28, {}, may-alias), {29}: (29, {}, may-alias), {30}: (30, {}, may-alias), {31}: (31, {}, may-alias), {32}: (32, {}, may-alias), {33}: (33, {}, may-alias), {34}: (34, {}, may-alias), {35}: (35, {}, may-alias), {36}: (36, {}, may-alias), {37}: (37, {}, may-alias), {38}: (38, {}, may-alias) }, entry_computation_layout={(s32[]{:T(128)}, f32[4096]{0:T(1024)}, f32[4096,4,14336]{2,1,0:T(4,128)}, f32[4096,4,14336]{2,1,0:T(4,128)}, f32[14336,4,4096]{2,1,0:T(4,128)}, /*index=5*/f32[4096,4]{0,1:T(4,128)}, f32[4096,4]{0,1:T(4,128)}, f32[4096,4,8,128]{3,2,1,0:T(8,128)}, f32[32,4,128,4096]{3,2,1,0:T(8,128)}, f32[4096,4,32,128]{3,2,1,0:T(8,128)}, /*index=10*/f32[4096,4,8,128]{3,2,1,0:T(8,128)}, f32[4096,128256]{1,0:T(8,128)}, f32[128256,4096]{1,0:T(8,128)}, s32[]{:T(128)}, f32[4096]{0:T(1024)}, /*index=15*/f32[4096,4,14336]{2,1,0:T(4,128)}, f32[4096,4,14336]{2,1,0:T(4,128)}, f32[14336,4,4096]{2,1,0:T(4,128)}, f32[4096,4]{0,1:T(4,128)}, f32[4096,4]{0,1:T(4,128)}, /*index=20*/f32[4096,4,8,128]{3,2,1,0:T(8,128)}, f32[32,4,128,4096]{3,2,1,0:T(8,128)}, f32[4096,4,32,128]{3,2,1,0:T(8,128)}, f32[4096,4,8,128]{3,2,1,0:T(8,128)}, f32[4096,128256]{1,0:T(8,128)}, /*index=25*/f32[128256,4096]{1,0:T(8,128)}, f32[4096]{0:T(1024)}, f32[4096,4,14336]{2,1,0:T(4,128)}, f32[4096,4,14336]{2,1,0:T(4,128)}, f32[14336,4,4096]{2,1,0:T(4,128)}, /*index=30*/f32[4096,4]{0,1:T(4,128)}, f32[4096,4]{0,1:T(4,128)}, f32[4096,4,8,128]{3,2,1,0:T(8,128)}, f32[32,4,128,4096]{3,2,1,0:T(8,128)}, f32[4096,4,32,128]{3,2,1,0:T(8,128)}, /*index=35*/f32[4096,4,8,128]{3,2,1,0:T(8,128)}, f32[4096,128256]{1,0:T(8,128)}, f32[128256,4096]{1,0:T(8,128)}, s32[]{:T(128)}, s32[4,128]{1,0:T(4,128)}, /*index=40*/s32[4,128]{1,0:T(4,128)}, s32[4,128]{1,0:T(4,128)}, s32[4,128]{1,0:T(4,128)}, s32[4,128]{1,0:T(4,128)})->(s32[]{:T(128)}, f32[4096]{0:T(1024)}, f32[4096,4,14336]{2,1,0:T(4,128)}, f32[4096,4,14336]{2,1,0:T(4,128)}, f32[14336,4,4096]{2,1,0:T(4,128)}, /*index=5*/f32[4096,4]{0,1:T(4,128)}, f32[4096,4]{0,1:T(4,128)}, f32[4096,4,8,128]{3,2,1,0:T(8,128)}, f32[32,4,128,4096]{3,2,1,0:T(8,128)}, f32[4096,4,32,128]{3,2,1,0:T(8,128)}, /*index=10*/f32[4096,4,8,128]{3,2,1,0:T(8,128)}, f32[4096,128256]{1,0:T(8,128)}, f32[128256,4096]{1,0:T(8,128)}, s32[]{:T(128)}, f32[4096]{0:T(1024)}, /*index=15*/f32[4096,4,14336]{2,1,0:T(4,128)}, f32[4096,4,14336]{2,1,0:T(4,128)}, f32[14336,4,4096]{2,1,0:T(4,128)}, f32[4096,4]{0,1:T(4,128)}, f32[4096,4]{0,1:T(4,128)}, /*index=20*/f32[4096,4,8,128]{3,2,1,0:T(8,128)}, f32[32,4,128,4096]{3,2,1,0:T(8,128)}, f32[4096,4,32,128]{3,2,1,0:T(8,128)}, f32[4096,4,8,128]{3,2,1,0:T(8,128)}, f32[4096,128256]{1,0:T(8,128)}, /*index=25*/f32[128256,4096]{1,0:T(8,128)}, f32[4096]{0:T(1024)}, f32[4096,4,14336]{2,1,0:T(4,128)}, f32[4096,4,14336]{2,1,0:T(4,128)}, f32[14336,4,4096]{2,1,0:T(4,128)}, /*index=30*/f32[4096,4]{0,1:T(4,128)}, f32[4096,4]{0,1:T(4,128)}, f32[4096,4,8,128]{3,2,1,0:T(8,128)}, f32[32,4,128,4096]{3,2,1,0:T(8,128)}, f32[4096,4,32,128]{3,2,1,0:T(8,128)}, /*index=35*/f32[4096,4,8,128]{3,2,1,0:T(8,128)}, f32[4096,128256]{1,0:T(8,128)}, f32[128256,4096]{1,0:T(8,128)}, s32[]{:T(128)}, f32[]{:T(128)}, /*index=40*/f32[]{:T(128)}, f32[]{:T(128)}, f32[]{:T(128)}, f32[]{:T(128)}, f32[]{:T(128)}, /*index=45*/f32[]{:T(128)}, f32[]{:T(128)}, f32[]{:T(128)}, s32[]{:T(128)}, f32[]{:T(128)})}, allow_spmd_sharding_propagation_to_parameters={false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false}, allow_spmd_sharding_propagation_to_output={false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,true,true,true,true,true,true,true,true,true,true,true}, num_partitions=4 +HloModule jit_train_step, is_scheduled=true, input_output_alias={ {0}: (0, {}, may-alias), {1}: (1, {}, may-alias), {2}: (2, {}, may-alias), {3}: (3, {}, may-alias), {4}: (4, {}, may-alias), {5}: (5, {}, may-alias), {6}: (6, {}, may-alias), {7}: (7, {}, may-alias), {8}: (8, {}, may-alias), {9}: (9, {}, may-alias), {10}: (10, {}, may-alias), {11}: (11, {}, may-alias), {12}: (12, {}, may-alias), {13}: (13, {}, may-alias), {14}: (14, {}, may-alias), {15}: (15, {}, may-alias), {16}: (16, {}, may-alias), {17}: (17, {}, may-alias), {18}: (18, {}, may-alias), {19}: (19, {}, may-alias), {20}: (20, {}, may-alias), {21}: (21, {}, may-alias), {22}: (22, {}, may-alias), {23}: (23, {}, may-alias), {24}: (24, {}, may-alias), {25}: (25, {}, may-alias), {26}: (26, {}, may-alias), {27}: (27, {}, may-alias), {28}: (28, {}, may-alias), {29}: (29, {}, may-alias), {30}: (30, {}, may-alias), {31}: (31, {}, may-alias), {32}: (32, {}, may-alias), {33}: (33, {}, may-alias), {34}: (34, {}, may-alias), {35}: (35, {}, may-alias), {36}: (36, {}, may-alias), {37}: (37, {}, may-alias), {38}: (38, {}, may-alias), {39}: (39, {}, may-alias), {40}: (40, {}, may-alias), {41}: (41, {}, may-alias), {42}: (42, {}, may-alias), {43}: (43, {}, may-alias), {44}: (44, {}, may-alias), {45}: (45, {}, may-alias), {46}: (46, {}, may-alias), {47}: (47, {}, may-alias), {48}: (48, {}, may-alias), {49}: (49, {}, may-alias), {50}: (50, {}, may-alias), {51}: (51, {}, may-alias), {52}: (52, {}, may-alias), {53}: (53, {}, may-alias), {54}: (54, {}, may-alias), {55}: (55, {}, may-alias), {56}: (56, {}, may-alias), {57}: (57, {}, may-alias), {58}: (58, {}, may-alias), {59}: (59, {}, may-alias), {60}: (60, {}, may-alias), {61}: (61, {}, may-alias), {62}: (62, {}, may-alias), {63}: (63, {}, may-alias), {64}: (64, {}, may-alias), {65}: (65, {}, may-alias), {66}: (66, {}, may-alias), {67}: (67, {}, may-alias), {68}: (68, {}, may-alias) }, entry_computation_layout={(f32[4096]{0:T(1024)}, u32[]{:T(128)}, u32[4]{0:T(128)}, u32[]{:T(128)}, u32[4]{0:T(128)}, /*index=5*/u32[]{:T(128)}, u32[4]{0:T(128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, /*index=10*/u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, /*index=15*/u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, f32[1024,4,14336]{2,1,0:T(4,128)}, /*index=20*/f32[1024,4,14336]{2,1,0:T(4,128)}, f32[14336,4,1024]{2,1,0:T(4,128)}, f32[4096,4]{0,1:T(4,128)}, f32[4096,4]{0,1:T(4,128)}, u32[4]{0:T(128)}, /*index=25*/u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, /*index=30*/f32[1024,4,8,128]{3,2,1,0:T(8,128)}, f32[32,4,128,1024]{3,2,1,0:T(8,128)}, f32[1024,4,32,128]{3,2,1,0:T(8,128)}, f32[1024,4,8,128]{3,2,1,0:T(8,128)}, f32[1024,128256]{1,0:T(8,128)}, /*index=35*/u32[]{:T(128)}, u32[4]{0:T(128)}, u32[]{:T(128)}, u32[4]{0:T(128)}, u32[]{:T(128)}, /*index=40*/u32[4]{0:T(128)}, f32[128256,1024]{1,0:T(8,128)}, s32[]{:T(128)}, f32[4096]{0:T(1024)}, f32[1024,4,14336]{2,1,0:T(4,128)}, /*index=45*/f32[1024,4,14336]{2,1,0:T(4,128)}, f32[14336,4,1024]{2,1,0:T(4,128)}, f32[4096,4]{0,1:T(4,128)}, f32[4096,4]{0,1:T(4,128)}, f32[1024,4,8,128]{3,2,1,0:T(8,128)}, /*index=50*/f32[32,4,128,1024]{3,2,1,0:T(8,128)}, f32[1024,4,32,128]{3,2,1,0:T(8,128)}, f32[1024,4,8,128]{3,2,1,0:T(8,128)}, f32[1024,128256]{1,0:T(8,128)}, f32[128256,1024]{1,0:T(8,128)}, /*index=55*/f32[4096]{0:T(1024)}, f32[1024,4,14336]{2,1,0:T(4,128)}, f32[1024,4,14336]{2,1,0:T(4,128)}, f32[14336,4,1024]{2,1,0:T(4,128)}, f32[4096,4]{0,1:T(4,128)}, /*index=60*/f32[4096,4]{0,1:T(4,128)}, f32[1024,4,8,128]{3,2,1,0:T(8,128)}, f32[32,4,128,1024]{3,2,1,0:T(8,128)}, f32[1024,4,32,128]{3,2,1,0:T(8,128)}, f32[1024,4,8,128]{3,2,1,0:T(8,128)}, /*index=65*/f32[1024,128256]{1,0:T(8,128)}, f32[128256,1024]{1,0:T(8,128)}, s32[]{:T(128)}, u32[]{:T(128)}, s32[1,128]{1,0:T(1,128)}, /*index=70*/s32[1,128]{1,0:T(1,128)}, s32[1,128]{1,0:T(1,128)}, s32[1,128]{1,0:T(1,128)}, s32[1,128]{1,0:T(1,128)})->(f32[4096]{0:T(1024)}, u32[]{:T(128)}, u32[4]{0:T(128)}, u32[]{:T(128)}, u32[4]{0:T(128)}, /*index=5*/u32[]{:T(128)}, u32[4]{0:T(128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, /*index=10*/u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, /*index=15*/u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, f32[1024,4,14336]{2,1,0:T(4,128)}, /*index=20*/f32[1024,4,14336]{2,1,0:T(4,128)}, f32[14336,4,1024]{2,1,0:T(4,128)}, f32[4096,4]{0,1:T(4,128)}, f32[4096,4]{0,1:T(4,128)}, u32[4]{0:T(128)}, /*index=25*/u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, /*index=30*/f32[1024,4,8,128]{3,2,1,0:T(8,128)}, f32[32,4,128,1024]{3,2,1,0:T(8,128)}, f32[1024,4,32,128]{3,2,1,0:T(8,128)}, f32[1024,4,8,128]{3,2,1,0:T(8,128)}, f32[1024,128256]{1,0:T(8,128)}, /*index=35*/u32[]{:T(128)}, u32[4]{0:T(128)}, u32[]{:T(128)}, u32[4]{0:T(128)}, u32[]{:T(128)}, /*index=40*/u32[4]{0:T(128)}, f32[128256,1024]{1,0:T(8,128)}, s32[]{:T(128)}, f32[4096]{0:T(1024)}, f32[1024,4,14336]{2,1,0:T(4,128)}, /*index=45*/f32[1024,4,14336]{2,1,0:T(4,128)}, f32[14336,4,1024]{2,1,0:T(4,128)}, f32[4096,4]{0,1:T(4,128)}, f32[4096,4]{0,1:T(4,128)}, f32[1024,4,8,128]{3,2,1,0:T(8,128)}, /*index=50*/f32[32,4,128,1024]{3,2,1,0:T(8,128)}, f32[1024,4,32,128]{3,2,1,0:T(8,128)}, f32[1024,4,8,128]{3,2,1,0:T(8,128)}, f32[1024,128256]{1,0:T(8,128)}, f32[128256,1024]{1,0:T(8,128)}, /*index=55*/f32[4096]{0:T(1024)}, f32[1024,4,14336]{2,1,0:T(4,128)}, f32[1024,4,14336]{2,1,0:T(4,128)}, f32[14336,4,1024]{2,1,0:T(4,128)}, f32[4096,4]{0,1:T(4,128)}, /*index=60*/f32[4096,4]{0,1:T(4,128)}, f32[1024,4,8,128]{3,2,1,0:T(8,128)}, f32[32,4,128,1024]{3,2,1,0:T(8,128)}, f32[1024,4,32,128]{3,2,1,0:T(8,128)}, f32[1024,4,8,128]{3,2,1,0:T(8,128)}, /*index=65*/f32[1024,128256]{1,0:T(8,128)}, f32[128256,1024]{1,0:T(8,128)}, s32[]{:T(128)}, u32[]{:T(128)}, f32[]{:T(128)}, /*index=70*/f32[]{:T(128)}, f32[]{:T(128)}, f32[]{:T(128)}, f32[]{:T(128)}, f32[]{:T(128)}, /*index=75*/f32[]{:T(128)}, f32[]{:T(128)}, f32[]{:T(128)}, s32[]{:T(128)}, f32[]{:T(128)})}, allow_spmd_sharding_propagation_to_parameters={false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false}, allow_spmd_sharding_propagation_to_output={false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,true,true,true,true,true,true,true,true,true,true,true}, num_partitions=4, frontend_attributes={arg_layout_modes="default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default",arg_memory_spaces="0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0",out_layout_modes="default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default",out_memory_spaces="0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0"} FileNames @@ -9,1992 +9,1992 @@ FileLocations StackFrames -%fused_computation (param_0.2: bf16[128256,4096], param_1.7: s32[1024]) -> bf16[512,4096] { - %param_0.2 = bf16[128256,4096]{1,0:T(8,128)(2,1)} parameter(0) - %param_1.7 = s32[1024]{0:T(1024)S(1)} parameter(1) - %custom-call.1 = s32[1024]{0:T(1024)} custom-call(%param_1.7), custom_call_target="AssumeGatherIndicesInBound", operand_layout_constraints={s32[1024]{0:T(1024)}}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} - %slice.6 = s32[512]{0:T(512)} slice(%custom-call.1), slice={[0:512]}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} - %reshape.342 = s32[4,128]{1,0:T(4,128)} reshape(%slice.6), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/broadcast_in_dim" stack_frame_id=0} - %transpose.326 = s32[4,128]{1,0:T(4,128)} transpose(%reshape.342), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/broadcast_in_dim" stack_frame_id=0} - %gather.4 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} gather(%param_0.2, %transpose.326), offset_dims={2}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=2, slice_sizes={1,4096}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} - %transpose.325 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} transpose(%gather.4), dimensions={0,1,2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} - ROOT %reshape.341 = bf16[512,4096]{1,0:T(8,128)(2,1)} reshape(%transpose.325), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} -} - -%region_33.38.clone (scatter-add.6: bf16[], scatter-add.7: bf16[]) -> bf16[] { - %scatter-add.6 = bf16[]{:T(256)} parameter(0), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/scatter-add"} - %scatter-add.7 = bf16[]{:T(256)} parameter(1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/scatter-add"} - ROOT %add.476 = bf16[]{:T(256)} add(%scatter-add.6, %scatter-add.7), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/add" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.1 (param_0.3: bf16[128256,4096], param_1.5: s32[512], param_2.4: bf16[512,4096]) -> bf16[128256,4096] { - %param_0.3 = bf16[128256,4096]{1,0:T(8,128)(2,1)} parameter(0) - %param_1.5 = s32[512]{0:T(512)S(1)} parameter(1) - %reshape.349 = s32[4,128]{1,0:T(4,128)} reshape(%param_1.5), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/broadcast_in_dim" stack_frame_id=0} - %transpose.331 = s32[4,128]{1,0:T(4,128)} transpose(%reshape.349), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/broadcast_in_dim" stack_frame_id=0} - %param_2.4 = bf16[512,4096]{1,0:T(8,128)(2,1)S(1)} parameter(2) - %reshape.350 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} reshape(%param_2.4), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while" stack_frame_id=0} - %transpose.332 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} transpose(%reshape.350), dimensions={0,1,2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while" stack_frame_id=0} - ROOT %scatter.2 = bf16[128256,4096]{1,0:T(8,128)(2,1)} scatter(%param_0.3, %transpose.331, %transpose.332), update_window_dims={2}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=2, to_apply=%region_33.38.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/scatter-add" stack_frame_id=0} -} - -%region_32.37 (reduce_sum.190: f32[], reduce_sum.191: f32[]) -> f32[] { - %reduce_sum.190 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.191 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.192 = f32[]{:T(128)} add(%reduce_sum.190, %reduce_sum.191), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.280.clone.clone.clone (param_0.1099: bf16[4,128,128256], param_1.1265: s32[4,128], param_2.1086: f32[4,128], param_3.785: f32[4,128], param_4.487: bf16[4,128], param_5.412: f32[4,128]) -> bf16[4,128,128256] { - %param_5.412 = f32[4,128]{1,0:T(4,128)S(1)} parameter(5) - %mul.1613 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_5.412), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - %param_3.785 = f32[4,128]{1,0:T(4,128)S(1)} parameter(3) - %mul.1612 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_3.785), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - %param_0.1099 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} parameter(0) - %convert_element_type.1044 = f32[4,128,128256]{2,1,0:T(8,128)} convert(%param_0.1099), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %param_4.487 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(4) - %sub.94 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_4.487), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %sub.93 = f32[4,128,128256]{2,1,0:T(8,128)} subtract(%convert_element_type.1044, %sub.94), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %exp.62 = f32[4,128,128256]{2,1,0:T(8,128)} exponential(%sub.93), metadata={op_name="jit(train_step)/jvp()/exp" stack_frame_id=0} - %mul.1611 = f32[4,128,128256]{2,1,0:T(8,128)} multiply(%mul.1612, %exp.62), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - %param_2.1086 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %div.823 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_2.1086), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} - %div.822 = f32[4,128,128256]{2,1,0:T(8,128)} divide(%mul.1611, %div.823), metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} - %param_1.1265 = s32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %eq.49 = s32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_1.1265), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} - %eq.48 = s32[4,128,128256]{2,1,0:T(8,128)} iota(), iota_dimension=2, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} - %eq.47 = pred[4,128,128256]{2,1,0:T(8,128)(4,1)} compare(%eq.49, %eq.48), direction=EQ, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} - %convert_element_type.1043 = f32[4,128,128256]{2,1,0:T(8,128)} convert(%eq.47), metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/convert_element_type" stack_frame_id=0} - %sub.92 = f32[4,128,128256]{2,1,0:T(8,128)} subtract(%div.822, %convert_element_type.1043), metadata={op_name="jit(train_step)/transpose(jvp())/sub" stack_frame_id=0} - %mul.1610 = f32[4,128,128256]{2,1,0:T(8,128)} multiply(%mul.1613, %sub.92), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - ROOT %convert_element_type.1042 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} convert(%mul.1610), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} -} - -%fused_computation.316.clone.clone (param_0.1100: f32[4,128], param_1.1266: bf16[4,128,4096], param_2.1088: bf16[4096]) -> bf16[4,128,4096] { - %param_2.1088 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(2) - %dot_general.387 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1088), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1266 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %convert_element_type.1046 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_1.1266), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %param_0.1100 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.1615 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_0.1100), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %mul.1614 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1046, %mul.1615), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %convert_element_type.1045 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1614), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - ROOT %dot_general.386 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.387, %convert_element_type.1045), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} -} - -%fused_computation.219 (param_0.1119: bf16[4,128,128256], param_1.1281: s32[4,128], param_2.1112: f32[4,128], param_3.801: f32[4,128], param_4.502: bf16[4,128], param_5.427: f32[4,128], param_6.299: f32[4,128], param_7.198: bf16[4,128,4096], param_8.116: bf16[4096]) -> (f32[], bf16[4096,128256,1]) { - %param_6.299 = f32[4,128]{1,0:T(4,128)S(1)} parameter(6) - %param_7.198 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(7) - %param_8.116 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(8) - %fusion.239.clone.1 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} fusion(%param_6.299, %param_7.198, %param_8.116), kind=kLoop, calls=%fused_computation.316.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %param_0.1119 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1281 = s32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %param_2.1112 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %param_3.801 = f32[4,128]{1,0:T(4,128)S(1)} parameter(3) - %param_4.502 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(4) - %param_5.427 = f32[4,128]{1,0:T(4,128)S(1)} parameter(5) - %multiply_convert_fusion.1.clone.1 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} fusion(%param_0.1119, %param_1.1281, %param_2.1112, %param_3.801, %param_4.502, /*index=5*/%param_5.427), kind=kLoop, calls=%fused_computation.280.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %convolution.88.clone.1 = bf16[4096,128256,1]{1,0,2:T(8,128)(2,1)} convolution(%fusion.239.clone.1, %multiply_convert_fusion.1.clone.1), window={size=4}, dim_labels=0fb_0io->bf0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/dot_general" stack_frame_id=0} - %bitcast.306 = bf16[4096,128256]{1,0:T(8,128)(2,1)} bitcast(%convolution.88.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/dot_general" stack_frame_id=0} - %convert_element_type.923 = f32[4096,128256]{1,0:T(8,128)} convert(%bitcast.306), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} - %square.157 = f32[4096,128256]{1,0:T(8,128)} multiply(%convert_element_type.923, %convert_element_type.923), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1006 = f32[]{:T(128)} constant(0) - %reduce.118 = f32[]{:T(128)} reduce(%square.157, %constant.1006), dimensions={0,1}, to_apply=%region_32.37, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.154 = (f32[]{:T(128)}, bf16[4096,128256,1]{1,0,2:T(8,128)(2,1)}) tuple(%reduce.118, %convolution.88.clone.1) -} - -%region_34.39 (reduce_sum.196: f32[], reduce_sum.197: f32[]) -> f32[] { - %reduce_sum.196 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.197 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.198 = f32[]{:T(128)} add(%reduce_sum.196, %reduce_sum.197), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.220 (param_0.1118: bf16[128256,4096]) -> f32[] { - %param_0.1118 = bf16[128256,4096]{1,0:T(8,128)(2,1)} parameter(0) - %convert_element_type.925 = f32[128256,4096]{1,0:T(8,128)} convert(%param_0.1118), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} - %square.159 = f32[128256,4096]{1,0:T(8,128)} multiply(%convert_element_type.925, %convert_element_type.925), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1005 = f32[]{:T(128)} constant(0) - ROOT %reduce.119 = f32[]{:T(128)} reduce(%square.159, %constant.1005), dimensions={0,1}, to_apply=%region_34.39, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} -} - -%region_60.65 (reduce_sum.338: f32[], reduce_sum.339: f32[]) -> f32[] { - %reduce_sum.338 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.339 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.329 = f32[]{:T(128)} add(%reduce_sum.338, %reduce_sum.339), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_46.51 (reduce_sum.259: f32[], reduce_sum.260: f32[]) -> f32[] { - %reduce_sum.259 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.260 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.261 = f32[]{:T(128)} add(%reduce_sum.259, %reduce_sum.260), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.221 (param_0.1106: f32[128256,4096], param_1.1269: f32[], param_2.1100: f32[], param_3.789: f32[], param_4.490: f32[128256,4096], param_5.415: f32[], param_6.287: bf16[128256,4096], param_7.186: pred[], param_8.104: f32[128256,4096]) -> (f32[], f32[128256,4096], f32[128256,4096], f32[128256,4096], f32[]) { - %param_0.1106 = f32[128256,4096]{1,0:T(8,128)} parameter(0) - %param_3.789 = f32[]{:T(128)S(6)} parameter(3) - %mul.1482.clone.1 = f32[128256,4096]{1,0:T(8,128)} broadcast(%param_3.789), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.186 = pred[]{:T(512)S(6)} parameter(7) - %select_n.242.clone.1 = pred[128256,4096]{1,0:T(8,128)(4,1)} broadcast(%param_7.186), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.287 = bf16[128256,4096]{1,0:T(8,128)(2,1)} parameter(6) - %convert_element_type.1017.clone.1 = f32[128256,4096]{1,0:T(8,128)} convert(%param_6.287), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} - %param_5.415 = f32[]{:T(128)} parameter(5) - %div.725.clone.1 = f32[128256,4096]{1,0:T(8,128)} broadcast(%param_5.415), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.724.clone.1 = f32[128256,4096]{1,0:T(8,128)} divide(%convert_element_type.1017.clone.1, %div.725.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.241.clone.1 = f32[128256,4096]{1,0:T(8,128)} select(%select_n.242.clone.1, %convert_element_type.1017.clone.1, %div.724.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %constant.907.clone.1 = f32[]{:T(128)} constant(0.1) - %broadcast.554.clone.1 = f32[128256,4096]{1,0:T(8,128)} broadcast(%constant.907.clone.1), dimensions={}, metadata={op_name="broadcast.61"} - %mul.1488.clone.1 = f32[128256,4096]{1,0:T(8,128)} multiply(%select_n.241.clone.1, %broadcast.554.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.104 = f32[128256,4096]{1,0:T(8,128)} parameter(8) - %constant.911.clone.1 = f32[]{:T(128)} constant(0.9) - %mul.1489.clone.1 = f32[128256,4096]{1,0:T(8,128)} broadcast(%constant.911.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1487.clone.1 = f32[128256,4096]{1,0:T(8,128)} multiply(%param_8.104, %mul.1489.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.776.clone.1 = f32[128256,4096]{1,0:T(8,128)} add(%mul.1488.clone.1, %mul.1487.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1100 = f32[]{:T(128)S(6)} parameter(2) - %div.721.clone.1 = f32[128256,4096]{1,0:T(8,128)} broadcast(%param_2.1100), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %integer_pow.60.clone.1 = f32[128256,4096]{1,0:T(8,128)} multiply(%select_n.241.clone.1, %select_n.241.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} - %constant.910.clone.1 = f32[]{:T(128)} constant(0.05) - %mul.1486.clone.1 = f32[128256,4096]{1,0:T(8,128)} broadcast(%constant.910.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1484.clone.1 = f32[128256,4096]{1,0:T(8,128)} multiply(%integer_pow.60.clone.1, %mul.1486.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.490 = f32[128256,4096]{1,0:T(8,128)} parameter(4) - %constant.909.clone.1 = f32[]{:T(128)} constant(0.95) - %mul.1485.clone.1 = f32[128256,4096]{1,0:T(8,128)} broadcast(%constant.909.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1483.clone.1 = f32[128256,4096]{1,0:T(8,128)} multiply(%param_4.490, %mul.1485.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.775.clone.1 = f32[128256,4096]{1,0:T(8,128)} add(%mul.1484.clone.1, %mul.1483.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1269 = f32[]{:T(128)S(6)} parameter(1) - %div.720.clone.1 = f32[128256,4096]{1,0:T(8,128)} broadcast(%param_1.1269), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.719.clone.1 = f32[128256,4096]{1,0:T(8,128)} divide(%add.775.clone.1, %div.720.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %sqrt.58.clone.1 = f32[128256,4096]{1,0:T(8,128)} sqrt(%div.719.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} - %constant.908.clone.1 = f32[]{:T(128)} constant(1e-08) - %add.774.clone.1 = f32[128256,4096]{1,0:T(8,128)} broadcast(%constant.908.clone.1), dimensions={}, metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %add.773.clone.1 = f32[128256,4096]{1,0:T(8,128)} add(%sqrt.58.clone.1, %add.774.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %multiply.256.clone.1 = f32[128256,4096]{1,0:T(8,128)} multiply(%div.721.clone.1, %add.773.clone.1), metadata={op_name="multiply.42"} - %div.718.clone.1 = f32[128256,4096]{1,0:T(8,128)} divide(%add.776.clone.1, %multiply.256.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1481.clone.1 = f32[128256,4096]{1,0:T(8,128)} multiply(%param_0.1106, %broadcast.554.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.772.clone.1 = f32[128256,4096]{1,0:T(8,128)} add(%div.718.clone.1, %mul.1481.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1480.clone.1 = f32[128256,4096]{1,0:T(8,128)} multiply(%mul.1482.clone.1, %add.772.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.771.clone.1 = f32[128256,4096]{1,0:T(8,128)} add(%param_0.1106, %mul.1480.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.160 = f32[128256,4096]{1,0:T(8,128)} multiply(%add.771.clone.1, %add.771.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.993 = f32[]{:T(128)} constant(0) - %reduce.120 = f32[]{:T(128)} reduce(%square.160, %constant.993), dimensions={0,1}, to_apply=%region_60.65, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.122.clone.1 = f32[]{:T(128)} reduce(%integer_pow.60.clone.1, %constant.993), dimensions={0,1}, to_apply=%region_46.51, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.135 = (f32[]{:T(128)}, f32[128256,4096]{1,0:T(8,128)}, f32[128256,4096]{1,0:T(8,128)}, f32[128256,4096]{1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.120, %add.771.clone.1, %add.775.clone.1, %add.776.clone.1, %reduce.122.clone.1) -} - -%region_59.64 (reduce_sum.331: f32[], reduce_sum.332: f32[]) -> f32[] { - %reduce_sum.331 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.332 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.323 = f32[]{:T(128)} add(%reduce_sum.331, %reduce_sum.332), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_45.50 (reduce_sum.253: f32[], reduce_sum.254: f32[]) -> f32[] { - %reduce_sum.253 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.254 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.255 = f32[]{:T(128)} add(%reduce_sum.253, %reduce_sum.254), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.222 (param_0.1107: f32[4096,128256], param_1.1270: f32[], param_2.1101: f32[], param_3.790: f32[], param_4.491: f32[4096,128256], param_5.416: f32[], param_6.288: bf16[4096,128256,1], param_7.187: pred[], param_8.105: f32[4096,128256]) -> (f32[], f32[4096,128256], f32[4096,128256], f32[4096,128256], f32[]) { - %param_0.1107 = f32[4096,128256]{1,0:T(8,128)} parameter(0) - %param_3.790 = f32[]{:T(128)S(6)} parameter(3) - %mul.1492.clone.1 = f32[4096,128256]{1,0:T(8,128)} broadcast(%param_3.790), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.187 = pred[]{:T(512)S(6)} parameter(7) - %select_n.246.clone.1 = pred[4096,128256]{1,0:T(8,128)(4,1)} broadcast(%param_7.187), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.288 = bf16[4096,128256,1]{1,0,2:T(8,128)(2,1)} parameter(6) - %bitcast.409.clone.1 = bf16[4096,128256]{1,0:T(8,128)(2,1)} bitcast(%param_6.288), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/dot_general" stack_frame_id=0} - %convert_element_type.1019.clone.1 = f32[4096,128256]{1,0:T(8,128)} convert(%bitcast.409.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} - %param_5.416 = f32[]{:T(128)} parameter(5) - %div.733.clone.1 = f32[4096,128256]{1,0:T(8,128)} broadcast(%param_5.416), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.732.clone.1 = f32[4096,128256]{1,0:T(8,128)} divide(%convert_element_type.1019.clone.1, %div.733.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.245.clone.1 = f32[4096,128256]{1,0:T(8,128)} select(%select_n.246.clone.1, %convert_element_type.1019.clone.1, %div.732.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %constant.913.clone.1 = f32[]{:T(128)} constant(0.1) - %broadcast.556.clone.1 = f32[4096,128256]{1,0:T(8,128)} broadcast(%constant.913.clone.1), dimensions={}, metadata={op_name="broadcast.62"} - %mul.1498.clone.1 = f32[4096,128256]{1,0:T(8,128)} multiply(%select_n.245.clone.1, %broadcast.556.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.105 = f32[4096,128256]{1,0:T(8,128)} parameter(8) - %constant.917.clone.1 = f32[]{:T(128)} constant(0.9) - %mul.1499.clone.1 = f32[4096,128256]{1,0:T(8,128)} broadcast(%constant.917.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1497.clone.1 = f32[4096,128256]{1,0:T(8,128)} multiply(%param_8.105, %mul.1499.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.782.clone.1 = f32[4096,128256]{1,0:T(8,128)} add(%mul.1498.clone.1, %mul.1497.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1101 = f32[]{:T(128)S(6)} parameter(2) - %div.729.clone.1 = f32[4096,128256]{1,0:T(8,128)} broadcast(%param_2.1101), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %integer_pow.61.clone.1 = f32[4096,128256]{1,0:T(8,128)} multiply(%select_n.245.clone.1, %select_n.245.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} - %constant.916.clone.1 = f32[]{:T(128)} constant(0.05) - %mul.1496.clone.1 = f32[4096,128256]{1,0:T(8,128)} broadcast(%constant.916.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1494.clone.1 = f32[4096,128256]{1,0:T(8,128)} multiply(%integer_pow.61.clone.1, %mul.1496.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.491 = f32[4096,128256]{1,0:T(8,128)} parameter(4) - %constant.915.clone.1 = f32[]{:T(128)} constant(0.95) - %mul.1495.clone.1 = f32[4096,128256]{1,0:T(8,128)} broadcast(%constant.915.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1493.clone.1 = f32[4096,128256]{1,0:T(8,128)} multiply(%param_4.491, %mul.1495.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.781.clone.1 = f32[4096,128256]{1,0:T(8,128)} add(%mul.1494.clone.1, %mul.1493.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1270 = f32[]{:T(128)S(6)} parameter(1) - %div.728.clone.1 = f32[4096,128256]{1,0:T(8,128)} broadcast(%param_1.1270), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.727.clone.1 = f32[4096,128256]{1,0:T(8,128)} divide(%add.781.clone.1, %div.728.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %sqrt.59.clone.1 = f32[4096,128256]{1,0:T(8,128)} sqrt(%div.727.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} - %constant.914.clone.1 = f32[]{:T(128)} constant(1e-08) - %add.780.clone.1 = f32[4096,128256]{1,0:T(8,128)} broadcast(%constant.914.clone.1), dimensions={}, metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %add.779.clone.1 = f32[4096,128256]{1,0:T(8,128)} add(%sqrt.59.clone.1, %add.780.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %multiply.257.clone.1 = f32[4096,128256]{1,0:T(8,128)} multiply(%div.729.clone.1, %add.779.clone.1), metadata={op_name="multiply.41"} - %div.726.clone.1 = f32[4096,128256]{1,0:T(8,128)} divide(%add.782.clone.1, %multiply.257.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1491.clone.1 = f32[4096,128256]{1,0:T(8,128)} multiply(%param_0.1107, %broadcast.556.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.778.clone.1 = f32[4096,128256]{1,0:T(8,128)} add(%div.726.clone.1, %mul.1491.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1490.clone.1 = f32[4096,128256]{1,0:T(8,128)} multiply(%mul.1492.clone.1, %add.778.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.777.clone.1 = f32[4096,128256]{1,0:T(8,128)} add(%param_0.1107, %mul.1490.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.161 = f32[4096,128256]{1,0:T(8,128)} multiply(%add.777.clone.1, %add.777.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.994 = f32[]{:T(128)} constant(0) - %reduce.121 = f32[]{:T(128)} reduce(%square.161, %constant.994), dimensions={0,1}, to_apply=%region_59.64, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.123.clone.1 = f32[]{:T(128)} reduce(%integer_pow.61.clone.1, %constant.994), dimensions={0,1}, to_apply=%region_45.50, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.136 = (f32[]{:T(128)}, f32[4096,128256]{1,0:T(8,128)}, f32[4096,128256]{1,0:T(8,128)}, f32[4096,128256]{1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.121, %add.777.clone.1, %add.781.clone.1, %add.782.clone.1, %reduce.123.clone.1) -} - -%region_25.30 (reduce_sum.154: f32[], reduce_sum.155: f32[]) -> f32[] { - %reduce_sum.154 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.155 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.156 = f32[]{:T(128)} add(%reduce_sum.154, %reduce_sum.155), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.239 (param_0.1124: f32[4,14336,4096]) -> f32[] { - %param_0.1124 = f32[4,14336,4096]{2,0,1:T(4,128)} parameter(0) - %bitcast.314 = f32[14336,4,4096]{2,1,0:T(4,128)} bitcast(%param_0.1124), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %square.164 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%bitcast.314, %bitcast.314), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1011 = f32[]{:T(128)} constant(0) - ROOT %reduce.124 = f32[]{:T(128)} reduce(%square.164, %constant.1011), dimensions={0,1,2}, to_apply=%region_25.30, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} -} - -%region_24.29 (reduce_sum.148: f32[], reduce_sum.149: f32[]) -> f32[] { - %reduce_sum.148 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.149 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.150 = f32[]{:T(128)} add(%reduce_sum.148, %reduce_sum.149), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_23.28 (reduce_sum.142: f32[], reduce_sum.143: f32[]) -> f32[] { - %reduce_sum.142 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.143 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.147 = f32[]{:T(128)} add(%reduce_sum.142, %reduce_sum.143), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.241 (param_0.1125: f32[4,4096,14336], param_1.1284: f32[4,4096,14336]) -> (f32[], f32[]) { - %param_0.1125 = f32[4,4096,14336]{2,0,1:T(4,128)} parameter(0) - %bitcast.318 = f32[4096,4,14336]{2,1,0:T(4,128)} bitcast(%param_0.1125), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %square.167 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%bitcast.318, %bitcast.318), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1012 = f32[]{:T(128)} constant(0) - %reduce.125 = f32[]{:T(128)} reduce(%square.167, %constant.1012), dimensions={0,1,2}, to_apply=%region_24.29, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %param_1.1284 = f32[4,4096,14336]{2,0,1:T(4,128)} parameter(1) - %bitcast.322.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} bitcast(%param_1.1284), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %square.170.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%bitcast.322.clone.1, %bitcast.322.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %reduce.126.clone.1 = f32[]{:T(128)} reduce(%square.170.clone.1, %constant.1012), dimensions={0,1,2}, to_apply=%region_23.28, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.155 = (f32[]{:T(128)}, f32[]{:T(128)}) tuple(%reduce.125, %reduce.126.clone.1) -} - -%fused_computation.244 (param_0.694: f32[14336,4,4096]) -> bf16[4,14336,4096] { - %param_0.694 = f32[14336,4,4096]{2,1,0:T(4,128)} parameter(0) - %copy.234 = bf16[14336,4,4096]{2,0,1:T(8,128)(2,1)} copy(%param_0.694), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'mlp\'][\'wo\'][\'kernel\']"} - ROOT %bitcast.323 = bf16[4,14336,4096]{2,1,0:T(8,128)(2,1)} bitcast(%copy.234), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} -} - -%fused_computation.245 (param_0.696: f32[4096,4,14336]) -> bf16[4,4096,14336] { - %param_0.696 = f32[4096,4,14336]{2,1,0:T(4,128)} parameter(0) - %copy.235 = bf16[4096,4,14336]{2,0,1:T(8,128)(2,1)} copy(%param_0.696), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'mlp\'][\'wi_1\'][\'kernel\']"} - ROOT %bitcast.324 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} bitcast(%copy.235), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} -} - -%fused_computation.246 (param_0.698: f32[4096,4,14336]) -> bf16[4,4096,14336] { - %param_0.698 = f32[4096,4,14336]{2,1,0:T(4,128)} parameter(0) - %copy.236 = bf16[4096,4,14336]{2,0,1:T(8,128)(2,1)} copy(%param_0.698), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'mlp\'][\'wi_0\'][\'kernel\']"} - ROOT %bitcast.325 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} bitcast(%copy.236), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} -} - -%region_52.57 (reduce_sum.289: f32[], reduce_sum.290: f32[]) -> f32[] { - %reduce_sum.289 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.290 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.294 = f32[]{:T(128)} add(%reduce_sum.289, %reduce_sum.290), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_38.43 (reduce_sum.217: f32[], reduce_sum.218: f32[]) -> f32[] { - %reduce_sum.217 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.218 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.219 = f32[]{:T(128)} add(%reduce_sum.217, %reduce_sum.218), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.247 (param_0.1114: f32[14336,4,4096], param_1.1277: f32[], param_2.1108: f32[], param_3.797: f32[], param_4.498: f32[14336,4,4096], param_5.423: f32[], param_6.295: f32[4,14336,4096], param_7.194: pred[], param_8.112: f32[14336,4,4096]) -> (f32[], f32[14336,4,4096], f32[14336,4,4096], f32[14336,4,4096], f32[]) { - %param_0.1114 = f32[14336,4,4096]{2,1,0:T(4,128)} parameter(0) - %param_3.797 = f32[]{:T(128)S(6)} parameter(3) - %mul.1550.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} broadcast(%param_3.797), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.194 = pred[]{:T(512)S(6)} parameter(7) - %select_n.274.clone.1 = pred[14336,4,4096]{2,1,0:T(4,128)(4,1)} broadcast(%param_7.194), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.295 = f32[4,14336,4096]{2,0,1:T(4,128)} parameter(6) - %bitcast.423.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} bitcast(%param_6.295), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.423 = f32[]{:T(128)} parameter(5) - %div.789.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} broadcast(%param_5.423), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.788.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} divide(%bitcast.423.clone.1, %div.789.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.273.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} select(%select_n.274.clone.1, %bitcast.423.clone.1, %div.788.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %constant.955.clone.1 = f32[]{:T(128)} constant(0.1) - %broadcast.586.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} broadcast(%constant.955.clone.1), dimensions={}, metadata={op_name="broadcast.69"} - %mul.1556.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%select_n.273.clone.1, %broadcast.586.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.112 = f32[14336,4,4096]{2,1,0:T(4,128)} parameter(8) - %constant.959.clone.1 = f32[]{:T(128)} constant(0.9) - %mul.1557.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} broadcast(%constant.959.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1555.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%param_8.112, %mul.1557.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.820.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} add(%mul.1556.clone.1, %mul.1555.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1108 = f32[]{:T(128)S(6)} parameter(2) - %div.785.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} broadcast(%param_2.1108), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %integer_pow.68.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%select_n.273.clone.1, %select_n.273.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} - %constant.958.clone.1 = f32[]{:T(128)} constant(0.05) - %mul.1554.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} broadcast(%constant.958.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1552.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%integer_pow.68.clone.1, %mul.1554.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.498 = f32[14336,4,4096]{2,1,0:T(4,128)} parameter(4) - %constant.957.clone.1 = f32[]{:T(128)} constant(0.95) - %mul.1553.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} broadcast(%constant.957.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1551.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%param_4.498, %mul.1553.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.819.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} add(%mul.1552.clone.1, %mul.1551.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1277 = f32[]{:T(128)S(6)} parameter(1) - %div.784.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} broadcast(%param_1.1277), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.783.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} divide(%add.819.clone.1, %div.784.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %sqrt.66.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} sqrt(%div.783.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} - %constant.956.clone.1 = f32[]{:T(128)} constant(1e-08) - %add.818.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} broadcast(%constant.956.clone.1), dimensions={}, metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %add.817.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} add(%sqrt.66.clone.1, %add.818.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %multiply.264.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%div.785.clone.1, %add.817.clone.1), metadata={op_name="multiply.34"} - %div.782.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} divide(%add.820.clone.1, %multiply.264.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1549.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%param_0.1114, %broadcast.586.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.816.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} add(%div.782.clone.1, %mul.1549.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1548.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%mul.1550.clone.1, %add.816.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.815.clone.1 = f32[14336,4,4096]{2,1,0:T(4,128)} add(%param_0.1114, %mul.1548.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.171 = f32[14336,4,4096]{2,1,0:T(4,128)} multiply(%add.815.clone.1, %add.815.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1001 = f32[]{:T(128)} constant(0) - %reduce.127 = f32[]{:T(128)} reduce(%square.171, %constant.1001), dimensions={0,1,2}, to_apply=%region_52.57, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.130.clone.1 = f32[]{:T(128)} reduce(%integer_pow.68.clone.1, %constant.1001), dimensions={0,1,2}, to_apply=%region_38.43, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.137 = (f32[]{:T(128)}, f32[14336,4,4096]{2,1,0:T(4,128)}, f32[14336,4,4096]{2,1,0:T(4,128)}, f32[14336,4,4096]{2,1,0:T(4,128)}, f32[]{:T(128)}) tuple(%reduce.127, %add.815.clone.1, %add.819.clone.1, %add.820.clone.1, %reduce.130.clone.1) -} - -%region_51.56 (reduce_sum.283: f32[], reduce_sum.287: f32[]) -> f32[] { - %reduce_sum.283 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.287 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.288 = f32[]{:T(128)} add(%reduce_sum.283, %reduce_sum.287), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_37.42 (reduce_sum.211: f32[], reduce_sum.212: f32[]) -> f32[] { - %reduce_sum.211 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.212 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.213 = f32[]{:T(128)} add(%reduce_sum.211, %reduce_sum.212), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.248 (param_0.1115: f32[4096,4,14336], param_1.1278: f32[], param_2.1109: f32[], param_3.798: f32[], param_4.499: f32[4096,4,14336], param_5.424: f32[], param_6.296: f32[4,4096,14336], param_7.195: pred[], param_8.113: f32[4096,4,14336]) -> (f32[], f32[4096,4,14336], f32[4096,4,14336], f32[4096,4,14336], f32[]) { - %param_0.1115 = f32[4096,4,14336]{2,1,0:T(4,128)} parameter(0) - %param_3.798 = f32[]{:T(128)S(6)} parameter(3) - %mul.1560.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%param_3.798), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.195 = pred[]{:T(512)S(6)} parameter(7) - %select_n.278.clone.1 = pred[4096,4,14336]{2,1,0:T(4,128)(4,1)} broadcast(%param_7.195), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.296 = f32[4,4096,14336]{2,0,1:T(4,128)} parameter(6) - %bitcast.425.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} bitcast(%param_6.296), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.424 = f32[]{:T(128)} parameter(5) - %div.797.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%param_5.424), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.796.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} divide(%bitcast.425.clone.1, %div.797.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.277.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} select(%select_n.278.clone.1, %bitcast.425.clone.1, %div.796.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %constant.961.clone.1 = f32[]{:T(128)} constant(0.1) - %broadcast.592.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%constant.961.clone.1), dimensions={}, metadata={op_name="broadcast.71"} - %mul.1564.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%select_n.277.clone.1, %broadcast.592.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.113 = f32[4096,4,14336]{2,1,0:T(4,128)} parameter(8) - %constant.965.clone.1 = f32[]{:T(128)} constant(0.9) - %broadcast.591.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%constant.965.clone.1), dimensions={}, metadata={op_name="broadcast.70"} - %mul.1563.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%param_8.113, %broadcast.591.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.825.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%mul.1564.clone.1, %mul.1563.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1109 = f32[]{:T(128)S(6)} parameter(2) - %div.793.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%param_2.1109), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %integer_pow.69.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%select_n.277.clone.1, %select_n.277.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} - %constant.964.clone.1 = f32[]{:T(128)} constant(0.05) - %broadcast.590.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%constant.964.clone.1), dimensions={}, metadata={op_name="broadcast.60"} - %mul.1562.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%integer_pow.69.clone.1, %broadcast.590.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.499 = f32[4096,4,14336]{2,1,0:T(4,128)} parameter(4) - %constant.963.clone.1 = f32[]{:T(128)} constant(0.95) - %broadcast.589.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%constant.963.clone.1), dimensions={}, metadata={op_name="broadcast.59"} - %mul.1561.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%param_4.499, %broadcast.589.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.824.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%mul.1562.clone.1, %mul.1561.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1278 = f32[]{:T(128)S(6)} parameter(1) - %div.792.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%param_1.1278), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.791.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} divide(%add.824.clone.1, %div.792.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %sqrt.67.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} sqrt(%div.791.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} - %constant.962.clone.1 = f32[]{:T(128)} constant(1e-08) - %broadcast.587.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%constant.962.clone.1), dimensions={}, metadata={op_name="broadcast.54"} - %add.823.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%sqrt.67.clone.1, %broadcast.587.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %multiply.265.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%div.793.clone.1, %add.823.clone.1), metadata={op_name="multiply.33"} - %div.790.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} divide(%add.825.clone.1, %multiply.265.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1559.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%param_0.1115, %broadcast.592.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.822.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%div.790.clone.1, %mul.1559.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1558.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%mul.1560.clone.1, %add.822.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.821.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%param_0.1115, %mul.1558.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.172 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%add.821.clone.1, %add.821.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1002 = f32[]{:T(128)} constant(0) - %reduce.128 = f32[]{:T(128)} reduce(%square.172, %constant.1002), dimensions={0,1,2}, to_apply=%region_51.56, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.131.clone.1 = f32[]{:T(128)} reduce(%integer_pow.69.clone.1, %constant.1002), dimensions={0,1,2}, to_apply=%region_37.42, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.138 = (f32[]{:T(128)}, f32[4096,4,14336]{2,1,0:T(4,128)}, f32[4096,4,14336]{2,1,0:T(4,128)}, f32[4096,4,14336]{2,1,0:T(4,128)}, f32[]{:T(128)}) tuple(%reduce.128, %add.821.clone.1, %add.824.clone.1, %add.825.clone.1, %reduce.131.clone.1) -} - -%region_50.55 (reduce_sum.280: f32[], reduce_sum.281: f32[]) -> f32[] { - %reduce_sum.280 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.281 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.282 = f32[]{:T(128)} add(%reduce_sum.280, %reduce_sum.281), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_36.41 (reduce_sum.205: f32[], reduce_sum.206: f32[]) -> f32[] { - %reduce_sum.205 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.206 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.210 = f32[]{:T(128)} add(%reduce_sum.205, %reduce_sum.206), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.249 (param_0.1116: f32[4096,4,14336], param_1.1279: f32[], param_2.1110: f32[], param_3.799: f32[], param_4.500: f32[4096,4,14336], param_5.425: f32[], param_6.297: f32[4,4096,14336], param_7.196: pred[], param_8.114: f32[4096,4,14336]) -> (f32[], f32[4096,4,14336], f32[4096,4,14336], f32[4096,4,14336], f32[]) { - %param_0.1116 = f32[4096,4,14336]{2,1,0:T(4,128)} parameter(0) - %param_3.799 = f32[]{:T(128)S(6)} parameter(3) - %mul.1567.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%param_3.799), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.196 = pred[]{:T(512)S(6)} parameter(7) - %select_n.282.clone.1 = pred[4096,4,14336]{2,1,0:T(4,128)(4,1)} broadcast(%param_7.196), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.297 = f32[4,4096,14336]{2,0,1:T(4,128)} parameter(6) - %bitcast.427.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} bitcast(%param_6.297), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.425 = f32[]{:T(128)} parameter(5) - %div.805.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%param_5.425), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.804.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} divide(%bitcast.427.clone.1, %div.805.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.281.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} select(%select_n.282.clone.1, %bitcast.427.clone.1, %div.804.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %constant.967.clone.1 = f32[]{:T(128)} constant(0.1) - %broadcast.598.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%constant.967.clone.1), dimensions={}, metadata={op_name="broadcast.71"} - %mul.1571.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%select_n.281.clone.1, %broadcast.598.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.114 = f32[4096,4,14336]{2,1,0:T(4,128)} parameter(8) - %constant.971.clone.1 = f32[]{:T(128)} constant(0.9) - %broadcast.597.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%constant.971.clone.1), dimensions={}, metadata={op_name="broadcast.70"} - %mul.1570.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%param_8.114, %broadcast.597.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.830.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%mul.1571.clone.1, %mul.1570.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1110 = f32[]{:T(128)S(6)} parameter(2) - %div.801.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%param_2.1110), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %integer_pow.70.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%select_n.281.clone.1, %select_n.281.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} - %constant.970.clone.1 = f32[]{:T(128)} constant(0.05) - %broadcast.596.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%constant.970.clone.1), dimensions={}, metadata={op_name="broadcast.60"} - %mul.1569.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%integer_pow.70.clone.1, %broadcast.596.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.500 = f32[4096,4,14336]{2,1,0:T(4,128)} parameter(4) - %constant.969.clone.1 = f32[]{:T(128)} constant(0.95) - %broadcast.595.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%constant.969.clone.1), dimensions={}, metadata={op_name="broadcast.59"} - %mul.1568.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%param_4.500, %broadcast.595.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.829.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%mul.1569.clone.1, %mul.1568.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1279 = f32[]{:T(128)S(6)} parameter(1) - %div.800.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%param_1.1279), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.799.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} divide(%add.829.clone.1, %div.800.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %sqrt.68.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} sqrt(%div.799.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} - %constant.968.clone.1 = f32[]{:T(128)} constant(1e-08) - %broadcast.593.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} broadcast(%constant.968.clone.1), dimensions={}, metadata={op_name="broadcast.54"} - %add.828.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%sqrt.68.clone.1, %broadcast.593.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %multiply.266.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%div.801.clone.1, %add.828.clone.1), metadata={op_name="multiply.32"} - %div.798.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} divide(%add.830.clone.1, %multiply.266.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1566.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%param_0.1116, %broadcast.598.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.827.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%div.798.clone.1, %mul.1566.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1565.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%mul.1567.clone.1, %add.827.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.826.clone.1 = f32[4096,4,14336]{2,1,0:T(4,128)} add(%param_0.1116, %mul.1565.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.173 = f32[4096,4,14336]{2,1,0:T(4,128)} multiply(%add.826.clone.1, %add.826.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1003 = f32[]{:T(128)} constant(0) - %reduce.129 = f32[]{:T(128)} reduce(%square.173, %constant.1003), dimensions={0,1,2}, to_apply=%region_50.55, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.132.clone.1 = f32[]{:T(128)} reduce(%integer_pow.70.clone.1, %constant.1003), dimensions={0,1,2}, to_apply=%region_36.41, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.139 = (f32[]{:T(128)}, f32[4096,4,14336]{2,1,0:T(4,128)}, f32[4096,4,14336]{2,1,0:T(4,128)}, f32[4096,4,14336]{2,1,0:T(4,128)}, f32[]{:T(128)}) tuple(%reduce.129, %add.826.clone.1, %add.829.clone.1, %add.830.clone.1, %reduce.132.clone.1) -} - -%region_30.35 (reduce_sum.178: f32[], reduce_sum.182: f32[]) -> f32[] { - %reduce_sum.178 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.182 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.183 = f32[]{:T(128)} add(%reduce_sum.178, %reduce_sum.182), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.267 (param_0.1120: f32[4,4096,32,128]) -> f32[] { - %param_0.1120 = f32[4,4096,32,128]{3,2,0,1:T(8,128)} parameter(0) - %bitcast.329 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} bitcast(%param_0.1120), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %square.176 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%bitcast.329, %bitcast.329), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1007 = f32[]{:T(128)} constant(0) - ROOT %reduce.133 = f32[]{:T(128)} reduce(%square.176, %constant.1007), dimensions={0,1,2,3}, to_apply=%region_30.35, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} -} - -%region_29.34 (reduce_sum.175: f32[], reduce_sum.176: f32[]) -> f32[] { - %reduce_sum.175 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.176 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.177 = f32[]{:T(128)} add(%reduce_sum.175, %reduce_sum.176), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.269 (param_0.1121: f32[4,32,128,4096]) -> f32[] { - %param_0.1121 = f32[4,32,128,4096]{3,2,0,1:T(8,128)} parameter(0) - %bitcast.333 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} bitcast(%param_0.1121), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %square.179 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%bitcast.333, %bitcast.333), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1008 = f32[]{:T(128)} constant(0) - ROOT %reduce.134 = f32[]{:T(128)} reduce(%square.179, %constant.1008), dimensions={0,1,2,3}, to_apply=%region_29.34, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} -} - -%fused_computation.270 (param_0.748: f32[32,4,128,4096]) -> bf16[4,32,128,4096] { - %param_0.748 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} parameter(0) - %copy.237 = bf16[32,4,128,4096]{3,2,0,1:T(8,128)(2,1)} copy(%param_0.748), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'self_attention\'][\'out\'][\'kernel\']"} - ROOT %bitcast.334 = bf16[4,32,128,4096]{3,2,1,0:T(8,128)(2,1)} bitcast(%copy.237), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} -} - -%region_57.62 (reduce_sum.317: f32[], reduce_sum.318: f32[]) -> f32[] { - %reduce_sum.317 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.318 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.316 = f32[]{:T(128)} add(%reduce_sum.317, %reduce_sum.318), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_43.48 (reduce_sum.241: f32[], reduce_sum.245: f32[]) -> f32[] { - %reduce_sum.241 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.245 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.246 = f32[]{:T(128)} add(%reduce_sum.241, %reduce_sum.245), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.271 (param_0.1109: f32[4096,4,32,128], param_1.1272: f32[], param_2.1103: f32[], param_3.792: f32[], param_4.493: f32[4096,4,32,128], param_5.418: f32[], param_6.290: f32[4,4096,32,128], param_7.189: pred[], param_8.107: f32[4096,4,32,128]) -> (f32[], f32[4096,4,32,128], f32[4096,4,32,128], f32[4096,4,32,128], f32[]) { - %param_0.1109 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} parameter(0) - %param_3.792 = f32[]{:T(128)S(6)} parameter(3) - %mul.1509.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} broadcast(%param_3.792), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.189 = pred[]{:T(512)S(6)} parameter(7) - %select_n.254.clone.1 = pred[4096,4,32,128]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.189), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.290 = f32[4,4096,32,128]{3,2,0,1:T(8,128)} parameter(6) - %bitcast.413.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} bitcast(%param_6.290), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.418 = f32[]{:T(128)} parameter(5) - %div.749.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} broadcast(%param_5.418), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.748.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} divide(%bitcast.413.clone.1, %div.749.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.253.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} select(%select_n.254.clone.1, %bitcast.413.clone.1, %div.748.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %constant.925.clone.1 = f32[]{:T(128)} constant(0.1) - %broadcast.564.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} broadcast(%constant.925.clone.1), dimensions={}, metadata={op_name="broadcast.63"} - %mul.1515.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%select_n.253.clone.1, %broadcast.564.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.107 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} parameter(8) - %constant.929.clone.1 = f32[]{:T(128)} constant(0.9) - %mul.1516.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} broadcast(%constant.929.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1514.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%param_8.107, %mul.1516.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.793.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} add(%mul.1515.clone.1, %mul.1514.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1103 = f32[]{:T(128)S(6)} parameter(2) - %div.745.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} broadcast(%param_2.1103), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %integer_pow.63.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%select_n.253.clone.1, %select_n.253.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} - %constant.928.clone.1 = f32[]{:T(128)} constant(0.05) - %mul.1513.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} broadcast(%constant.928.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1511.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%integer_pow.63.clone.1, %mul.1513.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.493 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} parameter(4) - %constant.927.clone.1 = f32[]{:T(128)} constant(0.95) - %mul.1512.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} broadcast(%constant.927.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1510.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%param_4.493, %mul.1512.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.792.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} add(%mul.1511.clone.1, %mul.1510.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1272 = f32[]{:T(128)S(6)} parameter(1) - %div.744.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} broadcast(%param_1.1272), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.743.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} divide(%add.792.clone.1, %div.744.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %sqrt.61.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} sqrt(%div.743.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} - %constant.926.clone.1 = f32[]{:T(128)} constant(1e-08) - %add.791.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} broadcast(%constant.926.clone.1), dimensions={}, metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %add.790.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} add(%sqrt.61.clone.1, %add.791.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %multiply.259.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%div.745.clone.1, %add.790.clone.1), metadata={op_name="multiply.39"} - %div.742.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} divide(%add.793.clone.1, %multiply.259.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1508.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%param_0.1109, %broadcast.564.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.789.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} add(%div.742.clone.1, %mul.1508.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1507.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%mul.1509.clone.1, %add.789.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.788.clone.1 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} add(%param_0.1109, %mul.1507.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.180 = f32[4096,4,32,128]{3,2,1,0:T(8,128)} multiply(%add.788.clone.1, %add.788.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.996 = f32[]{:T(128)} constant(0) - %reduce.135 = f32[]{:T(128)} reduce(%square.180, %constant.996), dimensions={0,1,2,3}, to_apply=%region_57.62, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.139.clone.1 = f32[]{:T(128)} reduce(%integer_pow.63.clone.1, %constant.996), dimensions={0,1,2,3}, to_apply=%region_43.48, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.140 = (f32[]{:T(128)}, f32[4096,4,32,128]{3,2,1,0:T(8,128)}, f32[4096,4,32,128]{3,2,1,0:T(8,128)}, f32[4096,4,32,128]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.135, %add.788.clone.1, %add.792.clone.1, %add.793.clone.1, %reduce.139.clone.1) -} - -%region_56.61 (reduce_sum.310: f32[], reduce_sum.311: f32[]) -> f32[] { - %reduce_sum.310 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.311 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.315 = f32[]{:T(128)} add(%reduce_sum.310, %reduce_sum.311), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_42.47 (reduce_sum.238: f32[], reduce_sum.239: f32[]) -> f32[] { - %reduce_sum.238 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.239 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.240 = f32[]{:T(128)} add(%reduce_sum.238, %reduce_sum.239), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.272 (param_0.1110: f32[32,4,128,4096], param_1.1273: f32[], param_2.1104: f32[], param_3.793: f32[], param_4.494: f32[32,4,128,4096], param_5.419: f32[], param_6.291: f32[4,32,128,4096], param_7.190: pred[], param_8.108: f32[32,4,128,4096]) -> (f32[], f32[32,4,128,4096], f32[32,4,128,4096], f32[32,4,128,4096], f32[]) { - %param_0.1110 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} parameter(0) - %param_3.793 = f32[]{:T(128)S(6)} parameter(3) - %mul.1519.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} broadcast(%param_3.793), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.190 = pred[]{:T(512)S(6)} parameter(7) - %select_n.258.clone.1 = pred[32,4,128,4096]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.190), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.291 = f32[4,32,128,4096]{3,2,0,1:T(8,128)} parameter(6) - %bitcast.415.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} bitcast(%param_6.291), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.419 = f32[]{:T(128)} parameter(5) - %div.757.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} broadcast(%param_5.419), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.756.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} divide(%bitcast.415.clone.1, %div.757.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.257.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} select(%select_n.258.clone.1, %bitcast.415.clone.1, %div.756.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %constant.931.clone.1 = f32[]{:T(128)} constant(0.1) - %broadcast.566.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} broadcast(%constant.931.clone.1), dimensions={}, metadata={op_name="broadcast.64"} - %mul.1525.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%select_n.257.clone.1, %broadcast.566.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.108 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} parameter(8) - %constant.935.clone.1 = f32[]{:T(128)} constant(0.9) - %mul.1526.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} broadcast(%constant.935.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1524.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%param_8.108, %mul.1526.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.799.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} add(%mul.1525.clone.1, %mul.1524.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1104 = f32[]{:T(128)S(6)} parameter(2) - %div.753.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} broadcast(%param_2.1104), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %integer_pow.64.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%select_n.257.clone.1, %select_n.257.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} - %constant.934.clone.1 = f32[]{:T(128)} constant(0.05) - %mul.1523.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} broadcast(%constant.934.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1521.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%integer_pow.64.clone.1, %mul.1523.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.494 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} parameter(4) - %constant.933.clone.1 = f32[]{:T(128)} constant(0.95) - %mul.1522.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} broadcast(%constant.933.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1520.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%param_4.494, %mul.1522.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.798.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} add(%mul.1521.clone.1, %mul.1520.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1273 = f32[]{:T(128)S(6)} parameter(1) - %div.752.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} broadcast(%param_1.1273), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.751.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} divide(%add.798.clone.1, %div.752.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %sqrt.62.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} sqrt(%div.751.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} - %constant.932.clone.1 = f32[]{:T(128)} constant(1e-08) - %add.797.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} broadcast(%constant.932.clone.1), dimensions={}, metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %add.796.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} add(%sqrt.62.clone.1, %add.797.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %multiply.260.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%div.753.clone.1, %add.796.clone.1), metadata={op_name="multiply.38"} - %div.750.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} divide(%add.799.clone.1, %multiply.260.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1518.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%param_0.1110, %broadcast.566.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.795.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} add(%div.750.clone.1, %mul.1518.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1517.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%mul.1519.clone.1, %add.795.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.794.clone.1 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} add(%param_0.1110, %mul.1517.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.181 = f32[32,4,128,4096]{3,2,1,0:T(8,128)} multiply(%add.794.clone.1, %add.794.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.997 = f32[]{:T(128)} constant(0) - %reduce.136 = f32[]{:T(128)} reduce(%square.181, %constant.997), dimensions={0,1,2,3}, to_apply=%region_56.61, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.140.clone.1 = f32[]{:T(128)} reduce(%integer_pow.64.clone.1, %constant.997), dimensions={0,1,2,3}, to_apply=%region_42.47, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.141 = (f32[]{:T(128)}, f32[32,4,128,4096]{3,2,1,0:T(8,128)}, f32[32,4,128,4096]{3,2,1,0:T(8,128)}, f32[32,4,128,4096]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.136, %add.794.clone.1, %add.798.clone.1, %add.799.clone.1, %reduce.140.clone.1) -} - -%region_47.52 (reduce_sum.262: f32[], reduce_sum.266: f32[]) -> f32[] { - %reduce_sum.262 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} - %reduce_sum.266 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} - ROOT %reduce_sum.267 = f32[]{:T(128)} add(%reduce_sum.262, %reduce_sum.266), metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +%fused_computation.408 (param_0.1130: s32[1,128]) -> s32[1,1,128] { + %param_0.1130 = s32[1,128]{1,0:T(1,128)S(1)} parameter(0) + %constant.378.clone.5 = s32[]{:T(128)} constant(0) + %broadcast.642 = s32[1,128]{1,0:T(1,128)} broadcast(%constant.378.clone.5), dimensions={}, metadata={op_name="broadcast.85"} + %lt.32 = pred[1,128]{1,0:T(4,128)(4,1)} compare(%param_0.1130, %broadcast.642), direction=LT, metadata={op_name="jit(train_step)/jvp()/lt" stack_frame_id=0} + %constant.390.clone.1 = s32[]{:T(128)} constant(128256) + %add.792 = s32[1,128]{1,0:T(1,128)} broadcast(%constant.390.clone.1), dimensions={}, metadata={op_name="jit(train_step)/jvp()/add" stack_frame_id=0} + %add.781 = s32[1,128]{1,0:T(1,128)} add(%param_0.1130, %add.792), metadata={op_name="jit(train_step)/jvp()/add" stack_frame_id=0} + %select_n.178 = s32[1,128]{1,0:T(1,128)} select(%lt.32, %add.781, %param_0.1130), metadata={op_name="jit(train_step)/jvp()/select_n" stack_frame_id=0} + ROOT %bitcast.560 = s32[1,1,128]{2,1,0:T(1,128)S(1)} bitcast(%select_n.178) } -%fused_computation.279 (param_0.1129: bf16[4,128,128256], param_1.1288: f32[4,128], param_2.1115: s32[4,128], param_3.803: bf16[4,128]) -> f32[4,128] { - %param_2.1115 = s32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %eq.30 = s32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_2.1115), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} - %eq.25 = s32[4,128,128256]{2,1,0:T(8,128)} iota(), iota_dimension=2, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} - %eq.24 = pred[4,128,128256]{2,1,0:T(8,128)(4,1)} compare(%eq.30, %eq.25), direction=EQ, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} - %param_0.1129 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} parameter(0) - %convert_element_type.950 = f32[4,128,128256]{2,1,0:T(8,128)} convert(%param_0.1129), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %param_3.803 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(3) - %sub.73 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_3.803), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %sub.64 = f32[4,128,128256]{2,1,0:T(8,128)} subtract(%convert_element_type.950, %sub.73), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %param_1.1288 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %sub.71 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_1.1288), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %sub.60 = f32[4,128,128256]{2,1,0:T(8,128)} subtract(%sub.64, %sub.71), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %constant.1017 = f32[]{:T(128)} constant(0) - %broadcast.511 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%constant.1017), dimensions={}, metadata={op_name="broadcast.83"} - %mul.1373 = f32[4,128,128256]{2,1,0:T(8,128)} select(%eq.24, %sub.60, %broadcast.511), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} - ROOT %reduce.137 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%mul.1373, %constant.1017), dimensions={2}, to_apply=%region_47.52, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} +%fused_computation.394 (param_0.1095: s32[512]) -> s32[1024] { + %constant.945 = s32[] constant(0), metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} + %broadcast.624 = s32[1024]{0:T(1024)} broadcast(%constant.945), dimensions={}, metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} + %param_0.1095 = s32[512]{0:T(512)S(1)} parameter(0) + %constant.946 = s32[] constant(2147483647), metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} + %pad.41 = s32[1024]{0:T(1024)} pad(%param_0.1095, %constant.946), padding=0_512, metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} + %constant.944 = s32[] constant(128255), metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} + %broadcast.623 = s32[1024]{0:T(1024)} broadcast(%constant.944), dimensions={}, metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} + ROOT %clamp.3 = s32[1024]{0:T(1024)S(1)} clamp(%broadcast.624, %pad.41, %broadcast.623), metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} } -%region_7.10 (reduce_sum.93: f32[], reduce_sum.94: f32[]) -> f32[] { - %reduce_sum.93 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} - %reduce_sum.94 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} - ROOT %reduce_sum.95 = f32[]{:T(128)} add(%reduce_sum.93, %reduce_sum.94), metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} +%fused_computation (param_0.2: bf16[128256,1024], param_1.7: s32[1024]) -> bf16[512,1024] { + %param_0.2 = bf16[128256,1024]{1,0:T(8,128)(2,1)} parameter(0) + %param_1.7 = s32[1024]{0:T(1024)S(1)} parameter(1) + %custom-call.1 = s32[1024]{0:T(1024)} custom-call(%param_1.7), custom_call_target="AssumeGatherIndicesInBound", operand_layout_constraints={s32[1024]{0:T(1024)}}, metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} + %slice.34 = s32[512]{0:T(512)} slice(%custom-call.1), slice={[0:512]}, metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} + %reshape.661 = s32[4,128]{1,0:T(4,128)} reshape(%slice.34), metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} + %transpose.326 = s32[4,128]{1,0:T(4,128)} transpose(%reshape.661), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} + %gather.4 = bf16[4,128,1024]{2,1,0:T(8,128)(2,1)} gather(%param_0.2, %transpose.326), offset_dims={2}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=2, slice_sizes={1,1024}, metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} + %transpose.325 = bf16[4,128,1024]{2,1,0:T(8,128)(2,1)} transpose(%gather.4), dimensions={0,1,2}, metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} + ROOT %reshape.660 = bf16[512,1024]{1,0:T(8,128)(2,1)S(1)} reshape(%transpose.325), metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} +} + +%fused_computation.421 (param_0.1265: f32[4096,4]) -> bf16[4,4096] { + %param_0.1265 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(0) + %bitcast.655 = f32[4,4096]{1,0:T(4,128)} bitcast(%param_0.1265), metadata={op_name="jit(train_step)/jvp()/transpose" stack_frame_id=0} + ROOT %convert.110 = bf16[4,4096]{1,0:T(4,128)(2,1)} convert(%bitcast.655) +} + +%fused_computation.420 (param_0.1264: f32[4096,4]) -> bf16[4,4096] { + %param_0.1264 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(0) + %bitcast.654 = f32[4,4096]{1,0:T(4,128)} bitcast(%param_0.1264), metadata={op_name="jit(train_step)/jvp()/transpose" stack_frame_id=0} + ROOT %convert.108 = bf16[4,4096]{1,0:T(4,128)(2,1)} convert(%bitcast.654) +} + +%fused_computation.416 (param_0.1146: s32[1,128]) -> (f32[1,128,1,1], f32[128]) { + %param_0.1146 = s32[1,128]{1,0:T(1,128)S(1)} parameter(0) + %convert_element_type.943 = f32[1,128]{1,0:T(1,128)} convert(%param_0.1146), metadata={op_name="jit(train_step)/convert_element_type" stack_frame_id=0} + %bitcast.570 = f32[1,128,1,1]{1,3,2,0:T(1,128)} bitcast(%convert_element_type.943), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %bitcast.571.clone.1 = f32[128]{0:T(128)S(1)} bitcast(%convert_element_type.943), metadata={op_name="jit(train_step)/convert_element_type" stack_frame_id=0} + ROOT %tuple.151 = (f32[1,128,1,1]{1,3,2,0:T(1,128)}, f32[128]{0:T(128)S(1)}) tuple(%bitcast.570, %bitcast.571.clone.1) +} + +%fused_computation.415 () -> f32[64] { + %constant.395.clone.1 = f32[]{:T(128)} constant(500000) + %broadcast.635 = f32[64]{0:T(128)} broadcast(%constant.395.clone.1), dimensions={}, metadata={op_name="broadcast.274"} + %iota.56 = s32[64]{0:T(128)} iota(), iota_dimension=0, metadata={op_name="jit(train_step)/iota" stack_frame_id=0} + %constant.396.clone.1 = s32[]{:T(128)} constant(2) + %broadcast.634 = s32[64]{0:T(128)} broadcast(%constant.396.clone.1), dimensions={}, metadata={op_name="broadcast.275"} + %mul.1510 = s32[64]{0:T(128)} multiply(%iota.56, %broadcast.634), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %convert_element_type.942 = f32[64]{0:T(128)} convert(%mul.1510), metadata={op_name="jit(train_step)/convert_element_type" stack_frame_id=0} + %constant.397.clone.1 = f32[]{:T(128)} constant(0.0078125) + %broadcast.633 = f32[64]{0:T(128)} broadcast(%constant.397.clone.1), dimensions={}, metadata={op_name="broadcast.276"} + %div.661 = f32[64]{0:T(128)} multiply(%convert_element_type.942, %broadcast.633), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + ROOT %pow.36 = f32[64]{0:T(128)S(1)} power(%broadcast.635, %div.661), metadata={op_name="jit(train_step)/pow" stack_frame_id=0} +} + +%fused_computation.369 (param_0.1070: f32[128], param_1.1137: f32[64]) -> (bf16[1,128,1,64], bf16[1,128,1,64]) { + %param_0.1070 = f32[128]{0:T(128)S(1)} parameter(0) + %div.631 = f32[1,128,1,64]{3,1,2,0:T(8,128)} broadcast(%param_0.1070), dimensions={1}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_1.1137 = f32[64]{0:T(128)S(1)} parameter(1) + %div.639 = f32[1,128,1,64]{3,1,2,0:T(8,128)} broadcast(%param_1.1137), dimensions={3}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.626 = f32[1,128,1,64]{3,1,2,0:T(8,128)} divide(%div.631, %div.639), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %sin.38 = f32[1,128,1,64]{3,1,2,0:T(8,128)} sine(%div.626), metadata={op_name="jit(train_step)/sin" stack_frame_id=0} + %convert_element_type.928 = bf16[1,128,1,64]{3,1,2,0:T(8,128)(2,1)S(1)} convert(%sin.38), metadata={op_name="jit(train_step)/convert_element_type" stack_frame_id=0} + %cos.41.clone.1 = f32[1,128,1,64]{3,1,2,0:T(8,128)} cosine(%div.626), metadata={op_name="jit(train_step)/cos" stack_frame_id=0} + %convert_element_type.927.clone.1 = bf16[1,128,1,64]{3,1,2,0:T(8,128)(2,1)S(1)} convert(%cos.41.clone.1), metadata={op_name="jit(train_step)/convert_element_type" stack_frame_id=0} + ROOT %tuple.148 = (bf16[1,128,1,64]{3,1,2,0:T(8,128)(2,1)S(1)}, bf16[1,128,1,64]{3,1,2,0:T(8,128)(2,1)S(1)}) tuple(%convert_element_type.928, %convert_element_type.927.clone.1) +} + +%fused_computation.384 (param_0.1065: bf16[1,128,1,64]) -> bf16[1,128,1,128] { + %param_0.1065 = bf16[1,128,1,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %constant.400.clone.2 = bf16[]{:T(256)} constant(-inf) + %pad.38 = bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1065, %constant.400.clone.2), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/concatenate" stack_frame_id=0} + %convert.133 = f32[1,128,1,128]{3,1,2,0:T(8,128)} convert(%pad.38) + %pad.37 = bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1065, %constant.400.clone.2), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/concatenate" stack_frame_id=0} + %convert.134 = f32[1,128,1,128]{3,1,2,0:T(8,128)} convert(%pad.37) + %maximum.38 = f32[1,128,1,128]{3,1,2,0:T(8,128)} maximum(%convert.133, %convert.134), metadata={op_name="jit(train_step)/concatenate" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","1","128"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + ROOT %convert.135 = bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)} convert(%maximum.38) +} -%fused_computation.284 (param_0.1130: bf16[4,128,128256], param_1.1289: bf16[4,128]) -> f32[4,128] { - %param_0.1130 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} parameter(0) - %convert_element_type.956 = f32[4,128,128256]{2,1,0:T(8,128)} convert(%param_0.1130), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %param_1.1289 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(1) - %sub.74 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_1.1289), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %sub.70 = f32[4,128,128256]{2,1,0:T(8,128)} subtract(%convert_element_type.956, %sub.74), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %exp.54 = f32[4,128,128256]{2,1,0:T(8,128)} exponential(%sub.70), metadata={op_name="jit(train_step)/jvp()/exp" stack_frame_id=0} - %constant.1018 = f32[]{:T(128)} constant(0) - ROOT %reduce.138 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%exp.54, %constant.1018), dimensions={2}, to_apply=%region_7.10, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} -} - -%region_31.36 (reduce_sum.184: f32[], reduce_sum.185: f32[]) -> f32[] { - %reduce_sum.184 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.185 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.189 = f32[]{:T(128)} add(%reduce_sum.184, %reduce_sum.185), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_28.33 (reduce_sum.169: f32[], reduce_sum.170: f32[]) -> f32[] { - %reduce_sum.169 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.170 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.171 = f32[]{:T(128)} add(%reduce_sum.169, %reduce_sum.170), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.290 (param_0.1122: f32[4,4096,8,128], param_1.1282: f32[4,4096,8,128]) -> (f32[], f32[]) { - %param_0.1122 = f32[4,4096,8,128]{3,2,0,1:T(8,128)} parameter(0) - %bitcast.350 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_0.1122), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %square.184 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%bitcast.350, %bitcast.350), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1009 = f32[]{:T(128)} constant(0) - %reduce.141 = f32[]{:T(128)} reduce(%square.184, %constant.1009), dimensions={0,1,2,3}, to_apply=%region_31.36, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %param_1.1282 = f32[4,4096,8,128]{3,2,0,1:T(8,128)} parameter(1) - %bitcast.354.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_1.1282), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %square.187.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%bitcast.354.clone.1, %bitcast.354.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %reduce.142.clone.1 = f32[]{:T(128)} reduce(%square.187.clone.1, %constant.1009), dimensions={0,1,2,3}, to_apply=%region_28.33, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.156 = (f32[]{:T(128)}, f32[]{:T(128)}) tuple(%reduce.141, %reduce.142.clone.1) -} - -%fused_computation.293 (param_0.807: f32[4096,4,8,128]) -> bf16[4,4096,8,128] { - %param_0.807 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} parameter(0) - %copy.238 = bf16[4096,4,8,128]{3,2,0,1:T(8,128)(2,1)} copy(%param_0.807), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'self_attention\'][\'value\'][\'kernel\']"} - ROOT %bitcast.355 = bf16[4,4096,8,128]{3,2,1,0:T(8,128)(2,1)S(1)} bitcast(%copy.238), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} -} - -%region_58.63 (reduce_sum.324: f32[], reduce_sum.325: f32[]) -> f32[] { - %reduce_sum.324 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.325 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.322 = f32[]{:T(128)} add(%reduce_sum.324, %reduce_sum.325), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_44.49 (reduce_sum.247: f32[], reduce_sum.248: f32[]) -> f32[] { - %reduce_sum.247 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.248 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.252 = f32[]{:T(128)} add(%reduce_sum.247, %reduce_sum.248), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.294 (param_0.1108: f32[4096,4,8,128], param_1.1271: f32[], param_2.1102: f32[], param_3.791: f32[], param_4.492: f32[4096,4,8,128], param_5.417: f32[], param_6.289: f32[4,4096,8,128], param_7.188: pred[], param_8.106: f32[4096,4,8,128]) -> (f32[], f32[4096,4,8,128], f32[4096,4,8,128], f32[4096,4,8,128], f32[]) { - %param_0.1108 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} parameter(0) - %param_3.791 = f32[]{:T(128)S(6)} parameter(3) - %mul.1502.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_3.791), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.188 = pred[]{:T(512)S(6)} parameter(7) - %select_n.250.clone.1 = pred[4096,4,8,128]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.188), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.289 = f32[4,4096,8,128]{3,2,0,1:T(8,128)} parameter(6) - %bitcast.411.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_6.289), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.417 = f32[]{:T(128)} parameter(5) - %div.741.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_5.417), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.740.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} divide(%bitcast.411.clone.1, %div.741.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.249.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} select(%select_n.250.clone.1, %bitcast.411.clone.1, %div.740.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %constant.919.clone.1 = f32[]{:T(128)} constant(0.1) - %broadcast.562.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.919.clone.1), dimensions={}, metadata={op_name="broadcast.66"} - %mul.1506.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%select_n.249.clone.1, %broadcast.562.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.106 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} parameter(8) - %constant.923.clone.1 = f32[]{:T(128)} constant(0.9) - %broadcast.561.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.923.clone.1), dimensions={}, metadata={op_name="broadcast.65"} - %mul.1505.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_8.106, %broadcast.561.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.787.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%mul.1506.clone.1, %mul.1505.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1102 = f32[]{:T(128)S(6)} parameter(2) - %div.737.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_2.1102), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %integer_pow.62.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%select_n.249.clone.1, %select_n.249.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} - %constant.922.clone.1 = f32[]{:T(128)} constant(0.05) - %broadcast.560.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.922.clone.1), dimensions={}, metadata={op_name="broadcast.56"} - %mul.1504.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%integer_pow.62.clone.1, %broadcast.560.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.492 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} parameter(4) - %constant.921.clone.1 = f32[]{:T(128)} constant(0.95) - %broadcast.559.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.921.clone.1), dimensions={}, metadata={op_name="broadcast.55"} - %mul.1503.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_4.492, %broadcast.559.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.786.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%mul.1504.clone.1, %mul.1503.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1271 = f32[]{:T(128)S(6)} parameter(1) - %div.736.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_1.1271), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.735.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} divide(%add.786.clone.1, %div.736.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %sqrt.60.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} sqrt(%div.735.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} - %constant.920.clone.1 = f32[]{:T(128)} constant(1e-08) - %broadcast.557.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.920.clone.1), dimensions={}, metadata={op_name="broadcast.52"} - %add.785.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%sqrt.60.clone.1, %broadcast.557.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %multiply.258.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%div.737.clone.1, %add.785.clone.1), metadata={op_name="multiply.40"} - %div.734.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} divide(%add.787.clone.1, %multiply.258.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1501.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_0.1108, %broadcast.562.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.784.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%div.734.clone.1, %mul.1501.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1500.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%mul.1502.clone.1, %add.784.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.783.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%param_0.1108, %mul.1500.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.188 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%add.783.clone.1, %add.783.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.995 = f32[]{:T(128)} constant(0) - %reduce.143 = f32[]{:T(128)} reduce(%square.188, %constant.995), dimensions={0,1,2,3}, to_apply=%region_58.63, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.145.clone.1 = f32[]{:T(128)} reduce(%integer_pow.62.clone.1, %constant.995), dimensions={0,1,2,3}, to_apply=%region_44.49, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.142 = (f32[]{:T(128)}, f32[4096,4,8,128]{3,2,1,0:T(8,128)}, f32[4096,4,8,128]{3,2,1,0:T(8,128)}, f32[4096,4,8,128]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.143, %add.783.clone.1, %add.786.clone.1, %add.787.clone.1, %reduce.145.clone.1) -} - -%region_55.60 (reduce_sum.304: f32[], reduce_sum.308: f32[]) -> f32[] { - %reduce_sum.304 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.308 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.309 = f32[]{:T(128)} add(%reduce_sum.304, %reduce_sum.308), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_41.46 (reduce_sum.232: f32[], reduce_sum.233: f32[]) -> f32[] { - %reduce_sum.232 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.233 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.234 = f32[]{:T(128)} add(%reduce_sum.232, %reduce_sum.233), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.295 (param_0.1111: f32[4096,4,8,128], param_1.1274: f32[], param_2.1105: f32[], param_3.794: f32[], param_4.495: f32[4096,4,8,128], param_5.420: f32[], param_6.292: f32[4,4096,8,128], param_7.191: pred[], param_8.109: f32[4096,4,8,128]) -> (f32[], f32[4096,4,8,128], f32[4096,4,8,128], f32[4096,4,8,128], f32[]) { - %param_0.1111 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} parameter(0) - %param_3.794 = f32[]{:T(128)S(6)} parameter(3) - %mul.1529.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_3.794), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.191 = pred[]{:T(512)S(6)} parameter(7) - %select_n.262.clone.1 = pred[4096,4,8,128]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.191), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.292 = f32[4,4096,8,128]{3,2,0,1:T(8,128)} parameter(6) - %bitcast.417.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_6.292), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.420 = f32[]{:T(128)} parameter(5) - %div.765.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_5.420), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.764.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} divide(%bitcast.417.clone.1, %div.765.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.261.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} select(%select_n.262.clone.1, %bitcast.417.clone.1, %div.764.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %constant.937.clone.1 = f32[]{:T(128)} constant(0.1) - %broadcast.572.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.937.clone.1), dimensions={}, metadata={op_name="broadcast.66"} - %mul.1533.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%select_n.261.clone.1, %broadcast.572.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.109 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} parameter(8) - %constant.941.clone.1 = f32[]{:T(128)} constant(0.9) - %broadcast.571.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.941.clone.1), dimensions={}, metadata={op_name="broadcast.65"} - %mul.1532.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_8.109, %broadcast.571.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.804.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%mul.1533.clone.1, %mul.1532.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1105 = f32[]{:T(128)S(6)} parameter(2) - %div.761.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_2.1105), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %integer_pow.65.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%select_n.261.clone.1, %select_n.261.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} - %constant.940.clone.1 = f32[]{:T(128)} constant(0.05) - %broadcast.570.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.940.clone.1), dimensions={}, metadata={op_name="broadcast.56"} - %mul.1531.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%integer_pow.65.clone.1, %broadcast.570.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.495 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} parameter(4) - %constant.939.clone.1 = f32[]{:T(128)} constant(0.95) - %broadcast.569.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.939.clone.1), dimensions={}, metadata={op_name="broadcast.55"} - %mul.1530.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_4.495, %broadcast.569.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.803.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%mul.1531.clone.1, %mul.1530.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1274 = f32[]{:T(128)S(6)} parameter(1) - %div.760.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_1.1274), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.759.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} divide(%add.803.clone.1, %div.760.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %sqrt.63.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} sqrt(%div.759.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} - %constant.938.clone.1 = f32[]{:T(128)} constant(1e-08) - %broadcast.567.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.938.clone.1), dimensions={}, metadata={op_name="broadcast.52"} - %add.802.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%sqrt.63.clone.1, %broadcast.567.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %multiply.261.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%div.761.clone.1, %add.802.clone.1), metadata={op_name="multiply.37"} - %div.758.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} divide(%add.804.clone.1, %multiply.261.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1528.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_0.1111, %broadcast.572.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.801.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%div.758.clone.1, %mul.1528.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1527.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%mul.1529.clone.1, %add.801.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.800.clone.1 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} add(%param_0.1111, %mul.1527.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.189 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} multiply(%add.800.clone.1, %add.800.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.998 = f32[]{:T(128)} constant(0) - %reduce.144 = f32[]{:T(128)} reduce(%square.189, %constant.998), dimensions={0,1,2,3}, to_apply=%region_55.60, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.146.clone.1 = f32[]{:T(128)} reduce(%integer_pow.65.clone.1, %constant.998), dimensions={0,1,2,3}, to_apply=%region_41.46, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.143 = (f32[]{:T(128)}, f32[4096,4,8,128]{3,2,1,0:T(8,128)}, f32[4096,4,8,128]{3,2,1,0:T(8,128)}, f32[4096,4,8,128]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.144, %add.800.clone.1, %add.803.clone.1, %add.804.clone.1, %reduce.146.clone.1) -} - -%fused_computation.311 (param_0.872: bf16[4,128,4096], param_1.941: f32[4,128], param_2.726: f32[4,128], param_3.452: bf16[4,128,4096], param_4.271: bf16[4096]) -> bf16[4,128,4096] { - %param_3.452 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) - %param_4.271 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(4) - %dot_general.375 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_4.271), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %dot_general.365 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%param_3.452, %dot_general.375), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %convert_element_type.973 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%dot_general.365), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %param_2.726 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %mul.1423 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_2.726), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %mul.1415 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.973, %mul.1423), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %param_0.872 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) - %convert_element_type.984 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_0.872), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %param_1.941 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %mul.1422 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_1.941), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %mul.1421 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.984, %mul.1422), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %add_any.138 = f32[4,128,4096]{2,1,0:T(8,128)} add(%mul.1415, %mul.1421), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/add_any" stack_frame_id=0} - ROOT %convert_element_type.971 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%add_any.138), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} +%fused_computation.385 (param_0.1067: bf16[1,128,1,64]) -> bf16[1,128,1,128] { + %param_0.1067 = bf16[1,128,1,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %constant.400.clone.1 = bf16[]{:T(256)} constant(-inf) + %pad.40 = bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1067, %constant.400.clone.1), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/concatenate" stack_frame_id=0} + %convert.136 = f32[1,128,1,128]{3,1,2,0:T(8,128)} convert(%pad.40) + %pad.39 = bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1067, %constant.400.clone.1), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/concatenate" stack_frame_id=0} + %convert.137 = f32[1,128,1,128]{3,1,2,0:T(8,128)} convert(%pad.39) + %maximum.39 = f32[1,128,1,128]{3,1,2,0:T(8,128)} maximum(%convert.136, %convert.137), metadata={op_name="jit(train_step)/concatenate" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","1","128"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + ROOT %convert.138 = bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)} convert(%maximum.39) +} + +%fused_computation.333 (param_0.933: f32[1024,4,8,128]) -> bf16[4,1024,8,128] { + %param_0.933 = f32[1024,4,8,128]{3,2,1,0:T(8,128)S(1)} parameter(0) + %copy.239 = bf16[1024,4,8,128]{3,0,2,1:T(8,128)(2,1)} copy(%param_0.933), sharding={devices=[4,1,1,1]<=[4]}, metadata={op_name="state[\'model\'][\'decoder\'][\'layers\'][\'self_attention\'][\'value\'][\'kernel\'].value"} + ROOT %bitcast.532 = bf16[4,1024,8,128]{3,1,2,0:T(8,128)(2,1)} bitcast(%copy.239), metadata={op_name="jit(train_step)/jvp()/transpose" stack_frame_id=0} +} + +%fused_computation.303 (param_0.855: f32[32,4,128,1024]) -> bf16[4,32,128,1024] { + %param_0.855 = f32[32,4,128,1024]{3,2,1,0:T(8,128)} parameter(0) + %copy.238 = bf16[32,4,128,1024]{3,2,0,1:T(8,128)(2,1)} copy(%param_0.855), sharding={devices=[1,1,1,4]<=[4]}, metadata={op_name="state[\'model\'][\'decoder\'][\'layers\'][\'self_attention\'][\'out\'][\'kernel\'].value"} + ROOT %bitcast.487 = bf16[4,32,128,1024]{3,2,1,0:T(8,128)(2,1)} bitcast(%copy.238), metadata={op_name="jit(train_step)/jvp()/transpose" stack_frame_id=0} +} + +%fused_computation.281 (param_0.811: f32[1024,4,14336]) -> bf16[4,1024,14336] { + %param_0.811 = f32[1024,4,14336]{2,1,0:T(4,128)} parameter(0) + %copy.237 = bf16[1024,4,14336]{2,0,1:T(8,128)(2,1)} copy(%param_0.811), sharding={devices=[4,1,1]<=[4]}, metadata={op_name="state[\'model\'][\'decoder\'][\'layers\'][\'mlp\'][\'wi_0\'][\'kernel\'].value"} + ROOT %bitcast.478 = bf16[4,1024,14336]{2,1,0:T(8,128)(2,1)} bitcast(%copy.237), metadata={op_name="jit(train_step)/jvp()/transpose" stack_frame_id=0} +} + +%fused_computation.280 (param_0.809: f32[1024,4,14336]) -> bf16[4,1024,14336] { + %param_0.809 = f32[1024,4,14336]{2,1,0:T(4,128)} parameter(0) + %copy.236 = bf16[1024,4,14336]{2,0,1:T(8,128)(2,1)} copy(%param_0.809), sharding={devices=[4,1,1]<=[4]}, metadata={op_name="state[\'model\'][\'decoder\'][\'layers\'][\'mlp\'][\'wi_1\'][\'kernel\'].value"} + ROOT %bitcast.477 = bf16[4,1024,14336]{2,1,0:T(8,128)(2,1)} bitcast(%copy.236), metadata={op_name="jit(train_step)/jvp()/transpose" stack_frame_id=0} +} + +%fused_computation.279 (param_0.807: f32[14336,4,1024]) -> bf16[4,14336,1024] { + %param_0.807 = f32[14336,4,1024]{2,1,0:T(4,128)} parameter(0) + %copy.235 = bf16[14336,4,1024]{2,0,1:T(8,128)(2,1)} copy(%param_0.807), sharding={devices=[1,1,4]<=[4]}, metadata={op_name="state[\'model\'][\'decoder\'][\'layers\'][\'mlp\'][\'wo\'][\'kernel\'].value"} + ROOT %bitcast.476 = bf16[4,14336,1024]{2,1,0:T(8,128)(2,1)} bitcast(%copy.235), metadata={op_name="jit(train_step)/jvp()/transpose" stack_frame_id=0} +} + +%fused_computation.54.clone.1 (param_0.1520: bf16[4,1,128,4096], param_1.1569: s32[], param_2.1198: bf16[1,128,4096]) -> bf16[4,1,128,4096] { + %param_0.1520 = bf16[4,1,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(0) + %param_2.1198 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %bitcast.881 = bf16[1,1,128,4096]{3,2,1,0:T(8,128)(2,1)} bitcast(%param_2.1198), metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_update_index_in_dim)/broadcast_in_dim" stack_frame_id=0} + %param_1.1569 = s32[]{:T(128)S(6)} parameter(1) + %constant.353.clone.24 = s32[]{:T(128)} constant(0) + ROOT %dynamic-update-slice.30 = bf16[4,1,128,4096]{3,2,1,0:T(8,128)(2,1)} dynamic-update-slice(%param_0.1520, %bitcast.881, %param_1.1569, %constant.353.clone.24, %constant.353.clone.24, /*index=5*/%constant.353.clone.24), metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_update_index_in_dim)/dynamic_update_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} +} + +%fused_computation.18.clone.1 (param_0.1540: bf16[4,32,128,1024], param_1.1582: s32[]) -> bf16[1,32,128,1024] { + %param_0.1540 = bf16[4,32,128,1024]{3,2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1582 = s32[]{:T(128)S(6)} parameter(1) + %constant.353.clone.29 = s32[]{:T(128)} constant(0) + ROOT %dynamic-slice.193 = bf16[1,32,128,1024]{3,2,1,0:T(8,128)(2,1)S(1)} dynamic-slice(%param_0.1540, %param_1.1582, %constant.353.clone.29, %constant.353.clone.29, %constant.353.clone.29), dynamic_slice_sizes={1,32,128,1024}, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_index_in_dim)/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} +} + +%fused_computation.19.clone.1 (param_0.1534: bf16[4,1024,32,128], param_1.1578: s32[]) -> bf16[1,1024,32,128] { + %param_0.1534 = bf16[4,1024,32,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) + %param_1.1578 = s32[]{:T(128)S(6)} parameter(1) + %constant.353.clone.28 = s32[]{:T(128)} constant(0) + ROOT %dynamic-slice.192 = bf16[1,1024,32,128]{1,3,2,0:T(8,128)(2,1)S(1)} dynamic-slice(%param_0.1534, %param_1.1578, %constant.353.clone.28, %constant.353.clone.28, %constant.353.clone.28), dynamic_slice_sizes={1,1024,32,128}, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_index_in_dim)/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} +} + +%region_2.3 (reduce_sum.73: f32[], reduce_sum.74: f32[]) -> f32[] { + %reduce_sum.73 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/reduce_sum"} + %reduce_sum.74 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/reduce_sum"} + ROOT %reduce_sum.78 = f32[]{:T(128)} add(%reduce_sum.73, %reduce_sum.74), metadata={op_name="reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%fused_computation.40.clone.1 (param_0.1522: bf16[1,128,4096]) -> f32[128] { + %param_0.1522 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.1051 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%param_0.1522), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %square.217 = f32[1,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1051, %convert_element_type.1051), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/square" stack_frame_id=0} + %constant.354.clone.9 = f32[]{:T(128)} constant(0) + ROOT %reduce.196 = f32[128]{0:T(128)S(1)} reduce(%square.217, %constant.354.clone.9), dimensions={0,2}, to_apply=%region_2.3, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/reduce_sum" stack_frame_id=0} +} + +%fused_computation.69.clone.1 (param_0.1523: f32[128]) -> f32[128] { + %param_0.1523 = f32[128]{0:T(128)S(1)} parameter(0) + %constant.355.clone.9 = f32[]{:T(128)} constant(0.000244140625) + %broadcast.721 = f32[128]{0:T(128)} broadcast(%constant.355.clone.9), dimensions={}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call" stack_frame_id=0} + %div.867 = f32[128]{0:T(128)} multiply(%param_0.1523, %broadcast.721), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/div" stack_frame_id=0} + %constant.356.clone.9 = f32[]{:T(128)} constant(1e-05) + %broadcast.720 = f32[128]{0:T(128)} broadcast(%constant.356.clone.9), dimensions={}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call" stack_frame_id=0} + %add.898 = f32[128]{0:T(128)} add(%div.867, %broadcast.720), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/add" stack_frame_id=0} + ROOT %rsqrt.118 = f32[128]{0:T(128)S(1)} rsqrt(%add.898), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/rsqrt" stack_frame_id=0} +} + +%convert_element_type.640.reduce_sub_computation (lhs.3: bf16[], rhs.3: bf16[]) -> bf16[] { + %lhs.3 = bf16[] parameter(0) + %rhs.3 = bf16[] parameter(1) + ROOT %add.662 = bf16[] add(%lhs.3, %rhs.3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%convert_element_type.638.reduce_sub_computation (lhs.2: bf16[], rhs.2: bf16[]) -> bf16[] { + %lhs.2 = bf16[] parameter(0) + %rhs.2 = bf16[] parameter(1) + ROOT %add.661 = bf16[] add(%lhs.2, %rhs.2), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%fused_computation.66.clone.1 (param_0.1521: bf16[4,4096], param_1.1570: s32[], param_2.1199: bf16[4,4096]) -> (bf16[4096], bf16[4096]) { + %param_0.1521 = bf16[4,4096]{1,0:T(4,128)(2,1)} parameter(0) + %param_1.1570 = s32[]{:T(128)S(6)} parameter(1) + %constant.353.clone.25 = s32[]{:T(128)} constant(0) + %dynamic_slice.291 = bf16[1,4096]{1,0:T(2,128)(2,1)} dynamic-slice(%param_0.1521, %param_1.1570, %constant.353.clone.25), dynamic_slice_sizes={1,4096}, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_index_in_dim)/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %constant.1016 = bf16[]{:T(256)} constant(-0), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %reduce.195 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} reduce(%dynamic_slice.291, %constant.1016), dimensions={0}, to_apply=%convert_element_type.640.reduce_sub_computation, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %param_2.1199 = bf16[4,4096]{1,0:T(4,128)(2,1)} parameter(2) + %dynamic_slice.283.clone.3 = bf16[1,4096]{1,0:T(2,128)(2,1)} dynamic-slice(%param_2.1199, %param_1.1570, %constant.353.clone.25), dynamic_slice_sizes={1,4096}, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_index_in_dim)/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %reduce.122.clone.3 = bf16[4096]{0:T(1024)(128)(2,1)} reduce(%dynamic_slice.283.clone.3, %constant.1016), dimensions={0}, to_apply=%convert_element_type.638.reduce_sub_computation, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + ROOT %tuple.199 = (bf16[4096]{0:T(1024)(128)(2,1)S(1)}, bf16[4096]{0:T(1024)(128)(2,1)}) tuple(%reduce.195, %reduce.122.clone.3) +} + +%fused_computation.48.clone.2.clone.1 (param_0.1536: bf16[1,128,4096], param_1.1579: f32[128], param_2.1205: bf16[4096]) -> bf16[128,4096,1] { + %param_2.1205 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(2) + %dot_general.472 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1205), dimensions={2}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/...k,k->...k/dot_general" stack_frame_id=0} + %convert.277 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%dot_general.472) + %param_0.1536 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.1057 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%param_0.1536), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %param_1.1579 = f32[128]{0:T(128)S(1)} parameter(1) + %mul.1808 = f32[1,128,4096]{2,1,0:T(8,128)} broadcast(%param_1.1579), dimensions={1}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %mul.1807 = f32[1,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1057, %mul.1808), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %convert_element_type.1056 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1807), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %convert.278 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%convert_element_type.1056) + %dot_general.471 = f32[1,128,4096]{2,1,0:T(8,128)} multiply(%convert.277, %convert.278), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/...k,k->...k/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","4096"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %convert.279 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} convert(%dot_general.471) + ROOT %bitcast.890 = bf16[128,4096,1]{1,0,2:T(8,128)(2,1)} bitcast(%convert.279) +} + +%fused_computation.17.clone.clone.clone.1 (param_0.1535: bf16[1,4096,32,128]) -> bf16[4096,32,128] { + %param_0.1535 = bf16[1,4096,32,128]{1,3,2,0:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.889 = bf16[4096,32,128]{0,2,1:T(8,128)(2,1)} bitcast(%param_0.1535), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} +} + +%fused_computation.57.clone.1 (param_0.1537: bf16[1,4096,32,128], param_1.1580: bf16[1,128,4096], param_2.1206: f32[128], param_3.812: bf16[4096]) -> bf16[1,128,32,128] { + %param_1.1580 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %param_2.1206 = f32[128]{0:T(128)S(1)} parameter(2) + %param_3.812 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(3) + %fusion.493 = bf16[128,4096,1]{1,0,2:T(8,128)(2,1)} fusion(%param_1.1580, %param_2.1206, %param_3.812), kind=kLoop, calls=%fused_computation.48.clone.2.clone.1 + %param_0.1537 = bf16[1,4096,32,128]{1,3,2,0:T(8,128)(2,1)S(1)} parameter(0) + %fusion.492 = bf16[4096,32,128]{0,2,1:T(8,128)(2,1)} fusion(%param_0.1537), kind=kLoop, calls=%fused_computation.17.clone.clone.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} + %convolution.160 = bf16[128,32,128]{2,0,1:T(8,128)(2,1)} convolution(%fusion.493, %fusion.492), window={size=32 pad=31_31 rhs_reversal=1}, dim_labels=bf0_i0o->b0f, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} + ROOT %bitcast.891 = bf16[1,128,32,128]{3,1,2,0:T(8,128)(2,1)S(1)} bitcast(%convolution.160), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} +} + +%fused_computation.49.clone.1 (param_0.1538: bf16[1,128,32,128]) -> (bf16[1,128,32,64], bf16[1,128,32,64]) { + %param_0.1538 = bf16[1,128,32,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %slice.82 = bf16[1,128,32,64]{3,1,2,0:T(8,128)(2,1)} slice(%param_0.1538), slice={[0:1], [0:128], [0:32], [64:128]}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/split" stack_frame_id=0} + %convert.290 = f32[1,128,32,64]{3,1,2,0:T(8,128)} convert(%slice.82) + %neg.139 = f32[1,128,32,64]{3,1,2,0:T(8,128)} negate(%convert.290), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/neg" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","32","64"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %convert.291 = bf16[1,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)} convert(%neg.139) + %slice.83 = bf16[1,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)} slice(%param_0.1538), slice={[0:1], [0:128], [0:32], [0:64]}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/split" stack_frame_id=0} + ROOT %tuple.201 = (bf16[1,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)}, bf16[1,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)}) tuple(%convert.291, %slice.83) +} + +%fused_computation.52.clone.1 (param_0.1539: bf16[1,128,32,64], param_1.1581: bf16[1,128,32,64], param_2.1207: bf16[1,128,32,128], param_3.813: bf16[128,128], param_4.500: bf16[128,128]) -> bf16[32,128,128] { + %param_2.1207 = bf16[1,128,32,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(2) + %convert.292 = f32[1,128,32,128]{3,1,2,0:T(8,128)} convert(%param_2.1207) + %param_4.500 = bf16[128,128]{1,0:T(8,128)(2,1)S(1)} parameter(4) + %mul.1812 = bf16[1,128,32,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_4.500), dimensions={1,3}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %convert.293 = f32[1,128,32,128]{3,1,2,0:T(8,128)} convert(%mul.1812) + %mul.1810 = f32[1,128,32,128]{3,1,2,0:T(8,128)} multiply(%convert.292, %convert.293), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","32","128"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %param_1.1581 = bf16[1,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) + %constant.1018 = bf16[]{:T(256)} constant(-inf) + %pad.73 = bf16[1,128,32,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_1.1581, %constant.1018), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/concatenate" stack_frame_id=0} + %convert.294 = f32[1,128,32,128]{3,1,2,0:T(8,128)} convert(%pad.73) + %param_0.1539 = bf16[1,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %pad.72 = bf16[1,128,32,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1539, %constant.1018), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/concatenate" stack_frame_id=0} + %convert.295 = f32[1,128,32,128]{3,1,2,0:T(8,128)} convert(%pad.72) + %maximum.55 = f32[1,128,32,128]{3,1,2,0:T(8,128)} maximum(%convert.294, %convert.295), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/concatenate" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","32","128"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %param_3.813 = bf16[128,128]{1,0:T(8,128)(2,1)S(1)} parameter(3) + %mul.1811 = bf16[1,128,32,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_3.813), dimensions={1,3}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %convert.296 = f32[1,128,32,128]{3,1,2,0:T(8,128)} convert(%mul.1811) + %mul.1809 = f32[1,128,32,128]{3,1,2,0:T(8,128)} multiply(%maximum.55, %convert.296), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","32","128"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %add.900 = f32[1,128,32,128]{3,1,2,0:T(8,128)} add(%mul.1810, %mul.1809), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/add" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","32","128"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %convert.297 = bf16[1,128,32,128]{3,1,2,0:T(8,128)(2,1)} convert(%add.900) + ROOT %bitcast.892 = bf16[32,128,128]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%convert.297), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/vmap(static_sparsity_splash)/jit(_splash_attention)/splash_mha_fwd_segmented_residuals/splash_mha_fwd_segmented_residuals/squeeze" stack_frame_id=0} +} + +%fused_computation.29.clone.1 (param_0.1528: bf16[4,1024,8,128], param_1.1574: s32[]) -> bf16[1,1024,8,128] { + %param_0.1528 = bf16[4,1024,8,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) + %param_1.1574 = s32[]{:T(128)S(6)} parameter(1) + %constant.353.clone.27 = s32[]{:T(128)} constant(0) + ROOT %dynamic-slice.191 = bf16[1,1024,8,128]{1,3,2,0:T(8,128)(2,1)S(1)} dynamic-slice(%param_0.1528, %param_1.1574, %constant.353.clone.27, %constant.353.clone.27, %constant.353.clone.27), dynamic_slice_sizes={1,1024,8,128}, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_index_in_dim)/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} +} + +%fused_computation.48.clone.1.clone.1 (param_0.1530: bf16[1,128,4096], param_1.1575: f32[128], param_2.1202: bf16[4096]) -> bf16[128,4096,1] { + %param_2.1202 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(2) + %dot_general.470 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1202), dimensions={2}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/...k,k->...k/dot_general" stack_frame_id=0} + %convert.274 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%dot_general.470) + %param_0.1530 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.1055 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%param_0.1530), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %param_1.1575 = f32[128]{0:T(128)S(1)} parameter(1) + %mul.1802 = f32[1,128,4096]{2,1,0:T(8,128)} broadcast(%param_1.1575), dimensions={1}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %mul.1801 = f32[1,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1055, %mul.1802), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %convert_element_type.1054 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1801), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %convert.275 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%convert_element_type.1054) + %dot_general.469 = f32[1,128,4096]{2,1,0:T(8,128)} multiply(%convert.274, %convert.275), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/...k,k->...k/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","4096"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %convert.276 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} convert(%dot_general.469) + ROOT %bitcast.886 = bf16[128,4096,1]{1,0,2:T(8,128)(2,1)} bitcast(%convert.276) +} + +%fused_computation.23.clone.clone.clone.1 (param_0.1529: bf16[1,4096,8,128]) -> bf16[4096,8,128] { + %param_0.1529 = bf16[1,4096,8,128]{1,3,2,0:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.885 = bf16[4096,8,128]{0,2,1:T(8,128)(2,1)} bitcast(%param_0.1529), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} +} + +%fused_computation.58.clone.1 (param_0.1531: bf16[1,4096,8,128], param_1.1576: bf16[1,128,4096], param_2.1203: f32[128], param_3.810: bf16[4096]) -> bf16[1,128,8,128] { + %param_1.1576 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %param_2.1203 = f32[128]{0:T(128)S(1)} parameter(2) + %param_3.810 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(3) + %fusion.491 = bf16[128,4096,1]{1,0,2:T(8,128)(2,1)} fusion(%param_1.1576, %param_2.1203, %param_3.810), kind=kLoop, calls=%fused_computation.48.clone.1.clone.1 + %param_0.1531 = bf16[1,4096,8,128]{1,3,2,0:T(8,128)(2,1)S(1)} parameter(0) + %fusion.490 = bf16[4096,8,128]{0,2,1:T(8,128)(2,1)} fusion(%param_0.1531), kind=kLoop, calls=%fused_computation.23.clone.clone.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} + %convolution.159 = bf16[128,8,128]{2,0,1:T(8,128)(2,1)} convolution(%fusion.491, %fusion.490), window={size=8 pad=7_7 rhs_reversal=1}, dim_labels=bf0_i0o->b0f, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} + ROOT %bitcast.887 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} bitcast(%convolution.159), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} +} + +%fused_computation.62.clone.1 (param_0.1532: bf16[1,128,8,128]) -> (bf16[1,128,8,64], bf16[1,128,8,64]) { + %param_0.1532 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %slice.80 = bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)} slice(%param_0.1532), slice={[0:1], [0:128], [0:8], [64:128]}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/split" stack_frame_id=0} + %convert.298 = f32[1,128,8,64]{3,1,2,0:T(8,128)} convert(%slice.80) + %neg.138 = f32[1,128,8,64]{3,1,2,0:T(8,128)} negate(%convert.298), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/neg" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","8","64"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %convert.299 = bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} convert(%neg.138) + %slice.81 = bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} slice(%param_0.1532), slice={[0:1], [0:128], [0:8], [0:64]}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/split" stack_frame_id=0} + ROOT %tuple.200 = (bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)}, bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)}) tuple(%convert.299, %slice.81) +} + +%fused_computation.65.clone.1 (param_0.1533: bf16[1,128,8,64], param_1.1577: bf16[1,128,8,64], param_2.1204: bf16[1,128,8,128], param_3.811: bf16[128,128], param_4.499: bf16[128,128]) -> bf16[8,128,128] { + %param_2.1204 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(2) + %convert.300 = f32[1,128,8,128]{3,1,2,0:T(8,128)} convert(%param_2.1204) + %param_4.499 = bf16[128,128]{1,0:T(8,128)(2,1)S(1)} parameter(4) + %mul.1806 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_4.499), dimensions={1,3}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %convert.301 = f32[1,128,8,128]{3,1,2,0:T(8,128)} convert(%mul.1806) + %mul.1804 = f32[1,128,8,128]{3,1,2,0:T(8,128)} multiply(%convert.300, %convert.301), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","8","128"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %param_1.1577 = bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) + %constant.1017 = bf16[]{:T(256)} constant(-inf) + %pad.71 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_1.1577, %constant.1017), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/concatenate" stack_frame_id=0} + %convert.302 = f32[1,128,8,128]{3,1,2,0:T(8,128)} convert(%pad.71) + %param_0.1533 = bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %pad.70 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1533, %constant.1017), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/concatenate" stack_frame_id=0} + %convert.303 = f32[1,128,8,128]{3,1,2,0:T(8,128)} convert(%pad.70) + %maximum.54 = f32[1,128,8,128]{3,1,2,0:T(8,128)} maximum(%convert.302, %convert.303), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/concatenate" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","8","128"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %param_3.811 = bf16[128,128]{1,0:T(8,128)(2,1)S(1)} parameter(3) + %mul.1805 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_3.811), dimensions={1,3}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %convert.304 = f32[1,128,8,128]{3,1,2,0:T(8,128)} convert(%mul.1805) + %mul.1803 = f32[1,128,8,128]{3,1,2,0:T(8,128)} multiply(%maximum.54, %convert.304), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","8","128"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %add.899 = f32[1,128,8,128]{3,1,2,0:T(8,128)} add(%mul.1804, %mul.1803), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/add" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","8","128"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %convert.305 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)} convert(%add.899) + ROOT %bitcast.888 = bf16[8,128,128]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%convert.305), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/vmap(static_sparsity_splash)/jit(_splash_attention)/splash_mha_fwd_segmented_residuals/splash_mha_fwd_segmented_residuals/squeeze" stack_frame_id=0} +} + +%fused_computation.28.clone.1 (param_0.1524: bf16[4,1024,8,128], param_1.1571: s32[]) -> bf16[1,1024,8,128] { + %param_0.1524 = bf16[4,1024,8,128]{3,1,2,0:T(8,128)(2,1)} parameter(0) + %param_1.1571 = s32[]{:T(128)S(6)} parameter(1) + %constant.353.clone.26 = s32[]{:T(128)} constant(0) + ROOT %dynamic-slice.190 = bf16[1,1024,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} dynamic-slice(%param_0.1524, %param_1.1571, %constant.353.clone.26, %constant.353.clone.26, %constant.353.clone.26), dynamic_slice_sizes={1,1024,8,128}, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_index_in_dim)/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} +} + +%fused_computation.48.clone.clone.1 (param_0.1526: bf16[1,128,4096], param_1.1572: f32[128], param_2.1200: bf16[4096]) -> bf16[128,4096,1] { + %param_2.1200 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(2) + %dot_general.468 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1200), dimensions={2}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/...k,k->...k/dot_general" stack_frame_id=0} + %convert.271 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%dot_general.468) + %param_0.1526 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.1053 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%param_0.1526), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %param_1.1572 = f32[128]{0:T(128)S(1)} parameter(1) + %mul.1800 = f32[1,128,4096]{2,1,0:T(8,128)} broadcast(%param_1.1572), dimensions={1}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %mul.1799 = f32[1,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1053, %mul.1800), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %convert_element_type.1052 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1799), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %convert.272 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%convert_element_type.1052) + %dot_general.467 = f32[1,128,4096]{2,1,0:T(8,128)} multiply(%convert.271, %convert.272), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/...k,k->...k/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","4096"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %convert.273 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} convert(%dot_general.467) + ROOT %bitcast.883 = bf16[128,4096,1]{1,0,2:T(8,128)(2,1)} bitcast(%convert.273) +} + +%fused_computation.21.clone.clone.clone.1 (param_0.1525: bf16[1,4096,8,128]) -> bf16[4096,8,128] { + %param_0.1525 = bf16[1,4096,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.882 = bf16[4096,8,128]{2,0,1:T(8,128)(2,1)} bitcast(%param_0.1525), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} +} + +%fused_computation.59.clone.1 (param_0.1527: bf16[1,4096,8,128], param_1.1573: bf16[1,128,4096], param_2.1201: f32[128], param_3.809: bf16[4096]) -> bf16[8,128,128] { + %param_1.1573 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %param_2.1201 = f32[128]{0:T(128)S(1)} parameter(2) + %param_3.809 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(3) + %fusion.489 = bf16[128,4096,1]{1,0,2:T(8,128)(2,1)} fusion(%param_1.1573, %param_2.1201, %param_3.809), kind=kLoop, calls=%fused_computation.48.clone.clone.1 + %param_0.1527 = bf16[1,4096,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %fusion.488 = bf16[4096,8,128]{2,0,1:T(8,128)(2,1)} fusion(%param_0.1527), kind=kLoop, calls=%fused_computation.21.clone.clone.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} + %convolution.158 = bf16[128,8,128]{2,0,1:T(8,128)(2,1)} convolution(%fusion.489, %fusion.488), window={size=8 pad=7_7 rhs_reversal=1}, dim_labels=bf0_i0o->b0f, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} + ROOT %bitcast.884 = bf16[8,128,128]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%convolution.158), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/vmap(static_sparsity_splash)/jit(_splash_attention)/splash_mha_fwd_segmented_residuals/splash_mha_fwd_segmented_residuals/squeeze" stack_frame_id=0} +} + +%fused_computation.53.clone.clone.clone.1 (param_0.1542: bf16[32,128,128]) -> bf16[128,32,128] { + %param_0.1542 = bf16[32,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.894 = bf16[128,32,128]{2,0,1:T(8,128)(2,1)} bitcast(%param_0.1542) +} + +%fused_computation.15.clone.clone.clone.clone.clone.clone.clone.1 (param_0.1541: bf16[1,32,128,4096]) -> bf16[32,128,4096] { + %param_0.1541 = bf16[1,32,128,4096]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.893 = bf16[32,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%param_0.1541), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} +} + +%region_3.5 (reduce_sum.79: f32[], reduce_sum.80: f32[]) -> f32[] { + %reduce_sum.79 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/reduce_sum"} + %reduce_sum.80 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/reduce_sum"} + ROOT %reduce_sum.81 = f32[]{:T(128)} add(%reduce_sum.79, %reduce_sum.80), metadata={op_name="reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%fused_computation.42.clone.1 (param_0.1543: bf16[1,128,4096], param_1.1583: bf16[1,32,128,4096], param_2.1208: bf16[32,128,128]) -> (f32[128], bf16[1,128,4096]) { + %param_0.1543 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert.306 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%param_0.1543) + %param_2.1208 = bf16[32,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %fusion.74.clone.3 = bf16[128,32,128]{2,0,1:T(8,128)(2,1)} fusion(%param_2.1208), kind=kLoop, calls=%fused_computation.53.clone.clone.clone.1 + %param_1.1583 = bf16[1,32,128,4096]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %fusion.59.clone.3 = bf16[32,128,4096]{2,1,0:T(8,128)(2,1)} fusion(%param_1.1583), kind=kLoop, calls=%fused_computation.15.clone.clone.clone.clone.clone.clone.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} + %convolution.58.clone.3 = bf16[128,4096,1]{1,0,2:T(8,128)(2,1)} convolution(%fusion.74.clone.3, %fusion.59.clone.3), window={size=32}, dim_labels=b0f_0io->bf0, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} + %bitcast.236.clone.3 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%convolution.58.clone.3), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} + %convert.307 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%bitcast.236.clone.3) + %add.671.clone.3 = f32[1,128,4096]{2,1,0:T(8,128)} add(%convert.306, %convert.307), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/add" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","4096"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %convert_element_type.1058 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%add.671.clone.3), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %square.218 = f32[1,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1058, %convert_element_type.1058), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/square" stack_frame_id=0} + %constant.354.clone.10 = f32[]{:T(128)} constant(0) + %reduce.197 = f32[128]{0:T(128)} reduce(%square.218, %constant.354.clone.10), dimensions={0,2}, to_apply=%region_3.5, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/reduce_sum" stack_frame_id=0} + %convert.308 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} convert(%add.671.clone.3) + ROOT %tuple.202 = (f32[128]{0:T(128)}, bf16[1,128,4096]{2,1,0:T(8,128)(2,1)}) tuple(%reduce.197, %convert.308) +} + +%fused_computation.11.clone.1 (param_0.1550: bf16[4,14336,1024], param_1.1588: s32[]) -> bf16[1,14336,1024] { + %param_0.1550 = bf16[4,14336,1024]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1588 = s32[]{:T(128)S(6)} parameter(1) + %constant.353.clone.32 = s32[]{:T(128)} constant(0) + ROOT %dynamic-slice.196 = bf16[1,14336,1024]{2,1,0:T(8,128)(2,1)S(1)} dynamic-slice(%param_0.1550, %param_1.1588, %constant.353.clone.32, %constant.353.clone.32), dynamic_slice_sizes={1,14336,1024}, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_index_in_dim)/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} +} + +%fused_computation.12.clone.1 (param_0.1549: bf16[4,1024,14336], param_1.1587: s32[]) -> bf16[1,1024,14336] { + %param_0.1549 = bf16[4,1024,14336]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1587 = s32[]{:T(128)S(6)} parameter(1) + %constant.353.clone.31 = s32[]{:T(128)} constant(0) + ROOT %dynamic-slice.195 = bf16[1,1024,14336]{2,1,0:T(8,128)(2,1)S(1)} dynamic-slice(%param_0.1549, %param_1.1587, %constant.353.clone.31, %constant.353.clone.31), dynamic_slice_sizes={1,1024,14336}, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_index_in_dim)/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} +} + +%fused_computation.13.clone.1 (param_0.1545: bf16[4,1024,14336], param_1.1584: s32[]) -> bf16[1,1024,14336] { + %param_0.1545 = bf16[4,1024,14336]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1584 = s32[]{:T(128)S(6)} parameter(1) + %constant.353.clone.30 = s32[]{:T(128)} constant(0) + ROOT %dynamic-slice.194 = bf16[1,1024,14336]{2,1,0:T(8,128)(2,1)S(1)} dynamic-slice(%param_0.1545, %param_1.1584, %constant.353.clone.30, %constant.353.clone.30), dynamic_slice_sizes={1,1024,14336}, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_index_in_dim)/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} +} + +%fused_computation.68.clone.1 (param_0.1544: f32[128]) -> f32[128] { + %param_0.1544 = f32[128]{0:T(128)S(1)} parameter(0) + %constant.355.clone.10 = f32[]{:T(128)} constant(0.000244140625) + %broadcast.723 = f32[128]{0:T(128)} broadcast(%constant.355.clone.10), dimensions={}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call" stack_frame_id=0} + %div.868 = f32[128]{0:T(128)} multiply(%param_0.1544, %broadcast.723), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/div" stack_frame_id=0} + %constant.356.clone.10 = f32[]{:T(128)} constant(1e-05) + %broadcast.722 = f32[128]{0:T(128)} broadcast(%constant.356.clone.10), dimensions={}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call" stack_frame_id=0} + %add.901 = f32[128]{0:T(128)} add(%div.868, %broadcast.722), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/add" stack_frame_id=0} + ROOT %rsqrt.119 = f32[128]{0:T(128)S(1)} rsqrt(%add.901), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/rsqrt" stack_frame_id=0} +} + +%fused_computation.46.clone.1.clone.1 (param_0.1547: bf16[1,128,4096], param_1.1585: f32[128], param_2.1209: bf16[4096]) -> bf16[128,4096] { + %param_2.1209 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(2) + %dot_general.474 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1209), dimensions={2}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/...k,k->...k/dot_general" stack_frame_id=0} + %convert.280 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%dot_general.474) + %param_0.1547 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.1060 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%param_0.1547), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %param_1.1585 = f32[128]{0:T(128)S(1)} parameter(1) + %mul.1814 = f32[1,128,4096]{2,1,0:T(8,128)} broadcast(%param_1.1585), dimensions={1}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %mul.1813 = f32[1,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1060, %mul.1814), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %convert_element_type.1059 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1813), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %convert.281 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%convert_element_type.1059) + %dot_general.473 = f32[1,128,4096]{2,1,0:T(8,128)} multiply(%convert.280, %convert.281), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/...k,k->...k/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","4096"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %convert.282 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} convert(%dot_general.473) + ROOT %bitcast.896 = bf16[128,4096]{1,0:T(8,128)(2,1)} bitcast(%convert.282) +} + +%fused_computation.10.clone.clone.clone.1 (param_0.1546: bf16[1,4096,14336]) -> bf16[4096,14336] { + %param_0.1546 = bf16[1,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) + ROOT %bitcast.895 = bf16[4096,14336]{1,0:T(8,128)(2,1)} bitcast(%param_0.1546), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} +} + +%fused_computation.24.clone.1 (param_0.1548: bf16[1,4096,14336], param_1.1586: bf16[1,128,4096], param_2.1210: f32[128], param_3.814: bf16[4096]) -> bf16[1,128,14336] { + %param_1.1586 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %param_2.1210 = f32[128]{0:T(128)S(1)} parameter(2) + %param_3.814 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(3) + %fusion.495 = bf16[128,4096]{1,0:T(8,128)(2,1)} fusion(%param_1.1586, %param_2.1210, %param_3.814), kind=kLoop, calls=%fused_computation.46.clone.1.clone.1 + %param_0.1548 = bf16[1,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) + %fusion.494 = bf16[4096,14336]{1,0:T(8,128)(2,1)} fusion(%param_0.1548), kind=kLoop, calls=%fused_computation.10.clone.clone.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} + %convolution.161 = bf16[128,14336]{1,0:T(8,128)(2,1)} convolution(%fusion.495, %fusion.494), dim_labels=bf_io->bf, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} + ROOT %bitcast.897 = bf16[1,128,14336]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%convolution.161), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} +} + +%fused_computation.46.clone.clone.clone.1 (param_0.1553: bf16[1,128,4096], param_1.1589: f32[128], param_2.1211: bf16[4096]) -> bf16[128,4096] { + %param_2.1211 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(2) + %dot_general.476 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1211), dimensions={2}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/...k,k->...k/dot_general" stack_frame_id=0} + %convert.283 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%dot_general.476) + %param_0.1553 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.1062 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%param_0.1553), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %param_1.1589 = f32[128]{0:T(128)S(1)} parameter(1) + %mul.1816 = f32[1,128,4096]{2,1,0:T(8,128)} broadcast(%param_1.1589), dimensions={1}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %mul.1815 = f32[1,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1062, %mul.1816), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %convert_element_type.1061 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1815), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %convert.284 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%convert_element_type.1061) + %dot_general.475 = f32[1,128,4096]{2,1,0:T(8,128)} multiply(%convert.283, %convert.284), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/...k,k->...k/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","4096"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %convert.285 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} convert(%dot_general.475) + ROOT %bitcast.900 = bf16[128,4096]{1,0:T(8,128)(2,1)} bitcast(%convert.285) +} + +%fused_computation.8.clone.clone.clone.clone.clone.clone.clone.clone.clone.clone.clone.clone.clone.1 (param_0.1552: bf16[1,4096,14336]) -> bf16[4096,14336] { + %param_0.1552 = bf16[1,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) + ROOT %bitcast.899 = bf16[4096,14336]{1,0:T(8,128)(2,1)} bitcast(%param_0.1552), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} +} + +%fused_computation.27.clone.clone.clone.clone.clone.clone.clone.1 (param_0.1554: bf16[1,4096,14336], param_1.1590: bf16[1,128,14336], param_2.1212: bf16[1,128,4096], param_3.815: f32[128], param_4.501: bf16[4096]) -> bf16[128,14336] { + %param_1.1590 = bf16[1,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %convert.286 = f32[1,128,14336]{2,1,0:T(8,128)} convert(%param_1.1590) + %constant.357.clone.10 = bf16[]{:T(256)} constant(1) + %jit_silu_.47 = bf16[1,128,14336]{2,1,0:T(8,128)(2,1)} broadcast(%constant.357.clone.10), dimensions={}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/jit(silu)" stack_frame_id=0} + %convert.287 = f32[1,128,14336]{2,1,0:T(8,128)} convert(%jit_silu_.47) + %neg.140 = f32[1,128,14336]{2,1,0:T(8,128)} negate(%convert.286), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/jit(silu)/neg" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","14336"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %exp.84 = f32[1,128,14336]{2,1,0:T(8,128)} exponential(%neg.140), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/jit(silu)/exp" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","14336"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %add.902 = f32[1,128,14336]{2,1,0:T(8,128)} add(%exp.84, %convert.287), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/jit(silu)/add" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","14336"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %div.869 = f32[1,128,14336]{2,1,0:T(8,128)} divide(%convert.287, %add.902), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/jit(silu)/div" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","14336"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %mul.1818 = f32[1,128,14336]{2,1,0:T(8,128)} multiply(%convert.286, %div.869), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/jit(silu)/mul" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","14336"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %param_2.1212 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %param_3.815 = f32[128]{0:T(128)S(1)} parameter(3) + %param_4.501 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(4) + %fusion.497 = bf16[128,4096]{1,0:T(8,128)(2,1)} fusion(%param_2.1212, %param_3.815, %param_4.501), kind=kLoop, calls=%fused_computation.46.clone.clone.clone.1 + %param_0.1554 = bf16[1,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) + %fusion.496 = bf16[4096,14336]{1,0:T(8,128)(2,1)} fusion(%param_0.1554), kind=kLoop, calls=%fused_computation.8.clone.clone.clone.clone.clone.clone.clone.clone.clone.clone.clone.clone.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} + %convolution.162 = bf16[128,14336]{1,0:T(8,128)(2,1)} convolution(%fusion.497, %fusion.496), dim_labels=bf_io->bf, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} + %bitcast.902 = bf16[1,128,14336]{2,1,0:T(8,128)(2,1)} bitcast(%convolution.162), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} + %convert.288 = f32[1,128,14336]{2,1,0:T(8,128)} convert(%bitcast.902) + %mul.1817 = f32[1,128,14336]{2,1,0:T(8,128)} multiply(%mul.1818, %convert.288), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","14336"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %convert.289 = bf16[1,128,14336]{2,1,0:T(8,128)(2,1)} convert(%mul.1817) + ROOT %bitcast.901 = bf16[128,14336]{1,0:T(8,128)(2,1)} bitcast(%convert.289) +} + +%fused_computation.6.clone.clone.clone.clone.clone.1 (param_0.1551: bf16[1,14336,4096]) -> bf16[14336,4096] { + %param_0.1551 = bf16[1,14336,4096]{2,1,0:T(8,128)(2,1)} parameter(0) + ROOT %bitcast.898 = bf16[14336,4096]{1,0:T(8,128)(2,1)} bitcast(%param_0.1551), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} +} + +%fused_computation.38.clone.1 (param_0.1555: bf16[1,128,4096], param_1.1591: bf16[1,14336,4096], param_2.1213: bf16[1,4096,14336], param_3.816: bf16[1,128,14336], param_4.502: f32[128], param_5.421: bf16[4096]) -> bf16[1,128,4096] { + %param_2.1213 = bf16[1,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(2) + %param_3.816 = bf16[1,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) + %param_0.1555 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %param_4.502 = f32[128]{0:T(128)S(1)} parameter(4) + %param_5.421 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(5) + %fusion.37.clone.7 = bf16[128,14336]{1,0:T(8,128)(2,1)} fusion(%param_2.1213, %param_3.816, %param_0.1555, %param_4.502, %param_5.421), kind=kOutput, calls=%fused_computation.27.clone.clone.clone.clone.clone.clone.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} + %param_1.1591 = bf16[1,14336,4096]{2,1,0:T(8,128)(2,1)} parameter(1) + %fusion.498 = bf16[14336,4096]{1,0:T(8,128)(2,1)} fusion(%param_1.1591), kind=kLoop, calls=%fused_computation.6.clone.clone.clone.clone.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} + %convolution.163 = bf16[128,4096]{1,0:T(8,128)(2,1)} convolution(%fusion.37.clone.7, %fusion.498), dim_labels=bf_io->bf, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} + %bitcast.903 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%convolution.163), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} + %convert.309 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%bitcast.903) + %convert.310 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%param_0.1555) + %add.903 = f32[1,128,4096]{2,1,0:T(8,128)} add(%convert.309, %convert.310), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/add" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","4096"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + ROOT %convert.311 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)S(1)} convert(%add.903) +} + +%async_computation.24 (param_0.1580: s8[33554432]) -> s8[4194304] { + %param_0.1580 = s8[33554432]{0:T(1024)(128)(4,1)} parameter(0) + ROOT %slice.133 = s8[4194304]{0:T(1024)(128)(4,1)S(1)} slice(%param_0.1580), slice={[0:4194304]}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%async_computation.25 (param_0.1581: s8[33554432]) -> s8[4194304] { + %param_0.1581 = s8[33554432]{0:T(1024)(128)(4,1)} parameter(0) + ROOT %slice.135 = s8[4194304]{0:T(1024)(128)(4,1)S(1)} slice(%param_0.1581), slice={[4194304:8388608]}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%async_computation.26 (param_0.1582: s8[33554432]) -> s8[4194304] { + %param_0.1582 = s8[33554432]{0:T(1024)(128)(4,1)} parameter(0) + ROOT %slice.137 = s8[4194304]{0:T(1024)(128)(4,1)S(1)} slice(%param_0.1582), slice={[8388608:12582912]}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%async_computation.27 (param_0.1583: s8[33554432]) -> s8[4194304] { + %param_0.1583 = s8[33554432]{0:T(1024)(128)(4,1)} parameter(0) + ROOT %slice.139 = s8[4194304]{0:T(1024)(128)(4,1)S(1)} slice(%param_0.1583), slice={[12582912:16777216]}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%async_computation.28 (param_0.1584: s8[33554432]) -> s8[4194304] { + %param_0.1584 = s8[33554432]{0:T(1024)(128)(4,1)} parameter(0) + ROOT %slice.141 = s8[4194304]{0:T(1024)(128)(4,1)S(1)} slice(%param_0.1584), slice={[16777216:20971520]}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%async_computation.29 (param_0.1585: s8[33554432]) -> s8[4194304] { + %param_0.1585 = s8[33554432]{0:T(1024)(128)(4,1)} parameter(0) + ROOT %slice.143 = s8[4194304]{0:T(1024)(128)(4,1)S(1)} slice(%param_0.1585), slice={[20971520:25165824]}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%async_computation.30 (param_0.1586: s8[33554432]) -> s8[4194304] { + %param_0.1586 = s8[33554432]{0:T(1024)(128)(4,1)} parameter(0) + ROOT %slice.145 = s8[4194304]{0:T(1024)(128)(4,1)S(1)} slice(%param_0.1586), slice={[25165824:29360128]}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%async_computation.31 (param_0.1587: s8[33554432]) -> s8[4194304] { + %param_0.1587 = s8[33554432]{0:T(1024)(128)(4,1)} parameter(0) + ROOT %slice.147 = s8[4194304]{0:T(1024)(128)(4,1)S(1)} slice(%param_0.1587), slice={[29360128:33554432]}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%wide.region_1.6_spmd.sunk.clone.clone.clone.sunk (wide.param.1: (s32[], bf16[1,128,4096], u32[4], u32[4,4], u32[4], /*index=5*/u32[4,4], u32[4], u32[4,4], u32[4], u32[4,4], /*index=10*/u32[4], u32[4,4], u32[4], u32[4,4], u32[4], /*index=15*/u32[4,4], u32[4], u32[4,4], u32[4], u32[4,4], /*index=20*/bf16[4,1,128,4096], bf16[4,4096], bf16[4,4096], bf16[4,1024,32,128], bf16[1,128,1,128], /*index=25*/bf16[1,128,1,128], bf16[4,1024,8,128], bf16[4,1024,8,128], bf16[4,32,128,1024], bf16[4,1024,14336], /*index=30*/bf16[4,1024,14336], bf16[4,14336,1024], u32[4], u32[4,4], u32[4], /*index=35*/u32[4,4], u32[4], u32[4,4], u32[4], u32[4,4], /*index=40*/u32[4], u32[4,4], u32[4], u32[4,4], u32[4], /*index=45*/u32[4,4], u32[4], u32[4,4], u32[4], u32[4,4], /*index=50*/s32[128], s32[], s8[1,1,1], s8[1,1,1], bf16[1,128,1,128], /*index=55*/bf16[1,128,1,128], s32[1,128])) -> (s32[], bf16[1,128,4096], u32[4], u32[4,4], u32[4], /*index=5*/u32[4,4], u32[4], u32[4,4], u32[4], u32[4,4], /*index=10*/u32[4], u32[4,4], u32[4], u32[4,4], u32[4], /*index=15*/u32[4,4], u32[4], u32[4,4], u32[4], u32[4,4], /*index=20*/bf16[4,1,128,4096], bf16[4,4096], bf16[4,4096], bf16[4,1024,32,128], bf16[1,128,1,128], /*index=25*/bf16[1,128,1,128], bf16[4,1024,8,128], bf16[4,1024,8,128], bf16[4,32,128,1024], bf16[4,1024,14336], /*index=30*/bf16[4,1024,14336], bf16[4,14336,1024], u32[4], u32[4,4], u32[4], /*index=35*/u32[4,4], u32[4], u32[4,4], u32[4], u32[4,4], /*index=40*/u32[4], u32[4,4], u32[4], u32[4,4], u32[4], /*index=45*/u32[4,4], u32[4], u32[4,4], u32[4], u32[4,4], /*index=50*/s32[128], s32[], s8[1,1,1], s8[1,1,1], bf16[1,128,1,128], /*index=55*/bf16[1,128,1,128], s32[1,128]) { + %constant.406.clone..sunk.1 = s32[]{:T(128)} constant(1) + %wide.param.1 = (s32[]{:T(128)}, bf16[1,128,4096]{2,1,0:T(8,128)(2,1)S(1)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, /*index=5*/u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, /*index=10*/u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, /*index=15*/u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, /*index=20*/bf16[4,1,128,4096]{3,2,1,0:T(8,128)(2,1)}, bf16[4,4096]{1,0:T(4,128)(2,1)}, bf16[4,4096]{1,0:T(4,128)(2,1)}, bf16[4,1024,32,128]{1,3,2,0:T(8,128)(2,1)}, bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)S(1)}, /*index=25*/bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)S(1)}, bf16[4,1024,8,128]{1,3,2,0:T(8,128)(2,1)}, bf16[4,1024,8,128]{3,1,2,0:T(8,128)(2,1)}, bf16[4,32,128,1024]{3,2,1,0:T(8,128)(2,1)}, bf16[4,1024,14336]{2,1,0:T(8,128)(2,1)}, /*index=30*/bf16[4,1024,14336]{2,1,0:T(8,128)(2,1)}, bf16[4,14336,1024]{2,1,0:T(8,128)(2,1)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, /*index=35*/u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, /*index=40*/u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, /*index=45*/u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, /*index=50*/s32[128]{0:T(128)S(1)}, s32[]{:T(128)}, s8[1,1,1]{2,1,0:T(4,128)(4,1)}, s8[1,1,1]{2,1,0:T(4,128)(4,1)}, bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)}, /*index=55*/bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)}, s32[1,128]{1,0:T(1,128)}) parameter(0) + %get-tuple-element.2248 = s32[]{:T(128)} get-tuple-element(%wide.param.1), index=0 + %copy.283 = s32[]{:T(128)S(6)} copy(%get-tuple-element.2248), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %add.904 = s32[]{:T(128)} add(%copy.283, %constant.406.clone..sunk.1), metadata={op_name="jit(train_step)/jvp()/while/body/add" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %get-tuple-element.2268 = bf16[4,1,128,4096]{3,2,1,0:T(8,128)(2,1)} get-tuple-element(%wide.param.1), index=20 + %get-tuple-element.2249 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)S(1)} get-tuple-element(%wide.param.1), index=1 + %bitcast_dynamic-update-slice_fusion.5 = bf16[4,1,128,4096]{3,2,1,0:T(8,128)(2,1)} fusion(%get-tuple-element.2268, %copy.283, %get-tuple-element.2249), kind=kLoop, calls=%fused_computation.54.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_update_index_in_dim)/dynamic_update_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[{"indices":["0","3"]}]}} + %get-tuple-element.2330 = bf16[4,32,128,1024]{3,2,1,0:T(8,128)(2,1)} get-tuple-element(%wide.param.1), index=28 + %dynamic-slice_convert_fusion.35 = bf16[1,32,128,1024]{3,2,1,0:T(8,128)(2,1)S(1)} fusion(%get-tuple-element.2330, %copy.283), kind=kLoop, calls=%fused_computation.18.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %all-gather.99 = bf16[1,32,128,4096]{3,2,1,0:T(8,128)(2,1)} all-gather(%dynamic-slice_convert_fusion.35), channel_id=103, replica_groups=mesh['axis_0'=1,'axis_1'=1,'axis_2'=4] {'axis_2'}, dimensions={3}, use_global_device_ids=true, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"barrier_config":{"barrier_type":"CUSTOM","id":"0"},"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %get-tuple-element.2354 = s8[1,1,1]{2,1,0:T(4,128)(4,1)} get-tuple-element(%wide.param.1), index=52, metadata={op_name="jit(train_step)/sharding_constraint" stack_frame_id=0} + %get-tuple-element.2355 = s8[1,1,1]{2,1,0:T(4,128)(4,1)} get-tuple-element(%wide.param.1), index=53, metadata={op_name="jit(train_step)/sharding_constraint" stack_frame_id=0} + %get-tuple-element.2325 = bf16[4,1024,32,128]{1,3,2,0:T(8,128)(2,1)} get-tuple-element(%wide.param.1), index=23 + %dynamic-slice_convert_fusion.36 = bf16[1,1024,32,128]{1,3,2,0:T(8,128)(2,1)S(1)} fusion(%get-tuple-element.2325, %copy.283), kind=kLoop, calls=%fused_computation.19.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %all-gather.100 = bf16[1,4096,32,128]{1,3,2,0:T(8,128)(2,1)S(1)} all-gather(%dynamic-slice_convert_fusion.36), channel_id=100, replica_groups=mesh['axis_0'=4,'axis_1'=1,'axis_2'=1] {'axis_0'}, dimensions={1}, use_global_device_ids=true, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"barrier_config":{"barrier_type":"CUSTOM","id":"0"},"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %fusion.499 = f32[128]{0:T(128)S(1)} fusion(%get-tuple-element.2249), kind=kLoop, calls=%fused_computation.40.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %add_rsqrt_fusion.4 = f32[128]{0:T(128)S(1)} fusion(%fusion.499), kind=kLoop, calls=%fused_computation.69.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/rsqrt" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %get-tuple-element.2324 = bf16[4,4096]{1,0:T(4,128)(2,1)} get-tuple-element(%wide.param.1), index=22 + %get-tuple-element.2323 = bf16[4,4096]{1,0:T(4,128)(2,1)} get-tuple-element(%wide.param.1), index=21 + %convert_reduce_fusion.9 = (bf16[4096]{0:T(1024)(128)(2,1)S(1)}, bf16[4096]{0:T(1024)(128)(2,1)}) fusion(%get-tuple-element.2324, %copy.283, %get-tuple-element.2323), kind=kLoop, calls=%fused_computation.66.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"fusion_formation_source":{"fusion_creating_compiler_pass":"TpuAdvancedMOF","fusion_id":"42"},"aliasing_operands":{"lists":[]}} + %get-tuple-element.2185 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} get-tuple-element(%convert_reduce_fusion.9), index=0, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %fusion.500 = bf16[1,128,32,128]{3,1,2,0:T(8,128)(2,1)S(1)} fusion(%all-gather.100, %get-tuple-element.2249, %add_rsqrt_fusion.4, %get-tuple-element.2185), kind=kOutput, calls=%fused_computation.57.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %bitcast.912 = s8[33554432]{0:T(1024)(128)(4,1)} bitcast(%all-gather.99) + %slice-start.24 = ((s8[33554432]{0:T(1024)(128)(4,1)}), s8[4194304]{0:T(1024)(128)(4,1)S(1)}, s32[]{:S(2)}) async-start(%bitcast.912), calls=%async_computation.24 + %slice_negate_fusion.12 = (bf16[1,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)}, bf16[1,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)}) fusion(%fusion.500), kind=kLoop, calls=%fused_computation.49.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/neg" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %get-tuple-element.2186 = bf16[1,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)} get-tuple-element(%slice_negate_fusion.12), index=1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/neg" stack_frame_id=0} + %get-tuple-element.2187 = bf16[1,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)} get-tuple-element(%slice_negate_fusion.12), index=0, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/neg" stack_frame_id=0} + %get-tuple-element.2327 = bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)S(1)} get-tuple-element(%wide.param.1), index=25 + %get-tuple-element.2326 = bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)S(1)} get-tuple-element(%wide.param.1), index=24 + %slice-start.25 = ((s8[33554432]{0:T(1024)(128)(4,1)}), s8[4194304]{0:T(1024)(128)(4,1)S(1)}, s32[]{:S(2)}) async-start(%bitcast.912), calls=%async_computation.25 + %bitcast.935 = bf16[128,128]{1,0:T(8,128)(2,1)S(1)} bitcast(%get-tuple-element.2327) + %bitcast.933 = bf16[128,128]{1,0:T(8,128)(2,1)S(1)} bitcast(%get-tuple-element.2326) + %fusion.501 = bf16[32,128,128]{2,1,0:T(8,128)(2,1)S(1)} fusion(%get-tuple-element.2186, %get-tuple-element.2187, %fusion.500, %bitcast.935, %bitcast.933), kind=kLoop, calls=%fused_computation.52.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/vmap(static_sparsity_splash)/jit(_splash_attention)/splash_mha_fwd_segmented_residuals/splash_mha_fwd_segmented_residuals/squeeze" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %get-tuple-element.2328 = bf16[4,1024,8,128]{1,3,2,0:T(8,128)(2,1)} get-tuple-element(%wide.param.1), index=26 + %dynamic-slice_convert_fusion.37 = bf16[1,1024,8,128]{1,3,2,0:T(8,128)(2,1)S(1)} fusion(%get-tuple-element.2328, %copy.283), kind=kLoop, calls=%fused_computation.29.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %slice-start.26 = ((s8[33554432]{0:T(1024)(128)(4,1)}), s8[4194304]{0:T(1024)(128)(4,1)S(1)}, s32[]{:S(2)}) async-start(%bitcast.912), calls=%async_computation.26 + %slice-start.27 = ((s8[33554432]{0:T(1024)(128)(4,1)}), s8[4194304]{0:T(1024)(128)(4,1)S(1)}, s32[]{:S(2)}) async-start(%bitcast.912), calls=%async_computation.27 + %all-gather.101 = bf16[1,4096,8,128]{1,3,2,0:T(8,128)(2,1)S(1)} all-gather(%dynamic-slice_convert_fusion.37), channel_id=101, replica_groups=mesh['axis_0'=4,'axis_1'=1,'axis_2'=1] {'axis_0'}, dimensions={1}, use_global_device_ids=true, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"barrier_config":{"barrier_type":"CUSTOM","id":"0"},"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %slice-start.28 = ((s8[33554432]{0:T(1024)(128)(4,1)}), s8[4194304]{0:T(1024)(128)(4,1)S(1)}, s32[]{:S(2)}) async-start(%bitcast.912), calls=%async_computation.28 + %fusion.502 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} fusion(%all-gather.101, %get-tuple-element.2249, %add_rsqrt_fusion.4, %get-tuple-element.2185), kind=kOutput, calls=%fused_computation.58.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %slice_negate_fusion.13 = (bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)}, bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)}) fusion(%fusion.502), kind=kLoop, calls=%fused_computation.62.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/neg" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %get-tuple-element.2188 = bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} get-tuple-element(%slice_negate_fusion.13), index=1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/neg" stack_frame_id=0} + %get-tuple-element.2189 = bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} get-tuple-element(%slice_negate_fusion.13), index=0, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/neg" stack_frame_id=0} + %bitcast.936 = bf16[128,128]{1,0:T(8,128)(2,1)S(1)} bitcast(%get-tuple-element.2327) + %bitcast.934 = bf16[128,128]{1,0:T(8,128)(2,1)S(1)} bitcast(%get-tuple-element.2326) + %fusion.503 = bf16[8,128,128]{2,1,0:T(8,128)(2,1)S(1)} fusion(%get-tuple-element.2188, %get-tuple-element.2189, %fusion.502, %bitcast.936, %bitcast.934), kind=kLoop, calls=%fused_computation.65.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/vmap(static_sparsity_splash)/jit(_splash_attention)/splash_mha_fwd_segmented_residuals/splash_mha_fwd_segmented_residuals/squeeze" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %get-tuple-element.2329 = bf16[4,1024,8,128]{3,1,2,0:T(8,128)(2,1)} get-tuple-element(%wide.param.1), index=27 + %slice-start.29 = ((s8[33554432]{0:T(1024)(128)(4,1)}), s8[4194304]{0:T(1024)(128)(4,1)S(1)}, s32[]{:S(2)}) async-start(%bitcast.912), calls=%async_computation.29 + %dynamic-slice_convert_fusion.38 = bf16[1,1024,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} fusion(%get-tuple-element.2329, %copy.283), kind=kLoop, calls=%fused_computation.28.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %slice-start.30 = ((s8[33554432]{0:T(1024)(128)(4,1)}), s8[4194304]{0:T(1024)(128)(4,1)S(1)}, s32[]{:S(2)}) async-start(%bitcast.912), calls=%async_computation.30 + %all-gather.102 = bf16[1,4096,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} all-gather(%dynamic-slice_convert_fusion.38), channel_id=102, replica_groups=mesh['axis_0'=4,'axis_1'=1,'axis_2'=1] {'axis_0'}, dimensions={1}, use_global_device_ids=true, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"barrier_config":{"barrier_type":"CUSTOM","id":"0"},"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %slice-start.31 = ((s8[33554432]{0:T(1024)(128)(4,1)}), s8[4194304]{0:T(1024)(128)(4,1)S(1)}, s32[]{:S(2)}) async-start(%bitcast.912), calls=%async_computation.31 + %fusion.504 = bf16[8,128,128]{2,1,0:T(8,128)(2,1)S(1)} fusion(%all-gather.102, %get-tuple-element.2249, %add_rsqrt_fusion.4, %get-tuple-element.2185), kind=kOutput, calls=%fused_computation.59.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %get-tuple-element.2352 = s32[128]{0:T(128)S(1)} get-tuple-element(%wide.param.1), index=50 + %squeeze.291 = s32[128,128]{1,0:T(8,128)S(1)} broadcast(%get-tuple-element.2352), dimensions={0}, metadata={op_name="jit(train_step)/shard_map/vmap(static_sparsity_splash)/jit(_splash_attention)/splash_mha_fwd_segmented_residuals/splash_mha_fwd_segmented_residuals/squeeze" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"output_chunk_bound_config":{"output_chunk_bound":["8"]},"aliasing_operands":{"lists":[]}} + %squeeze.292 = s32[8,128]{1,0:T(8,128)S(1)} broadcast(%get-tuple-element.2352), dimensions={1}, metadata={op_name="jit(train_step)/shard_map/vmap(static_sparsity_splash)/jit(_splash_attention)/splash_mha_fwd_segmented_residuals/splash_mha_fwd_segmented_residuals/squeeze" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"output_chunk_bound_config":{"output_chunk_bound":["128"]},"aliasing_operands":{"lists":[]}} + %iota.68 = s32[128,128]{1,0:T(8,128)S(1)} iota(), iota_dimension=0, metadata={op_name="jit(train_step)/shard_map/vmap(static_sparsity_splash)/jit(_splash_attention)/broadcast_in_dim" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %splash_mha_fwd_segmented_residuals.3 = (f32[128,128]{1,0:T(8,128)}, f32[128,128]{1,0:T(8,128)}, f32[128,128]{1,0:T(8,128)}, bf16[32,128,128]{2,1,0:T(8,128)(2,1)S(1)}, f32[32,128,128]{2,1,0:T(8,128)}) custom-call(%get-tuple-element.2354, %get-tuple-element.2355, %fusion.501, %fusion.503, %fusion.504, /*index=5*/%squeeze.291, %squeeze.292, %iota.68), custom_call_target="tpu_custom_call", operand_layout_constraints={s8[1,1,1]{2,1,0}, s8[1,1,1]{2,1,0}, bf16[32,128,128]{2,1,0}, bf16[8,128,128]{2,1,0}, bf16[8,128,128]{2,1,0}, s32[128,128]{1,0}, s32[8,128]{1,0}, s32[128,128]{1,0}}, frontend_attributes={kernel_metadata={ +"xprof_metadata":"{\"block_q\": 128, \"block_kv\": 128, \"block_kv_compute\": 128, \"block_q_dkv\": 128, \"block_kv_dkv\": 128, \"block_kv_dkv_compute\": 128, \"block_q_dq\": 128, \"block_kv_dq\": 128, \"use_fused_bwd_kernel\": false, \"q_layout\": 1, \"k_layout\": 1, \"v_layout\": 1}" +}}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/vmap(static_sparsity_splash)/jit(_splash_attention)/splash_mha_fwd_segmented_residuals/splash_mha_fwd_segmented_residuals/pallas_call" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"custom_call_config":{"body":"TUzvUgFNTElSMjMuMC4wZ2l0AAFVCwEDBQcJAQMLAz8NDxETFRcZGx0fISMlJykrLS8xMzU3OTs9P0FDRUdJA7IECgQzAfUHFxMXFxMPFxcbCwsXCwsLEycPFwsXExNzUwsXGxMXCxMTF6WFCxMXEwsXDwsLCwsLCwsLCwsLCwsXEwsLJwsLCxMLFxcLCwsLFxcLCzMXFxcXEwsPSwsXCwtzCw8LCwsLNxsLGwsbCxsLGwtTGwsjCyMLIwsbCxsLGwURYYGRjWFNKgIqAgH5GxsbGxsbGxsbGw8PEwsLKwsLEwsTCw8TCw8LEwsTCxMLCxMLDxMLFx8LEwsLEwtTFwsTCwsTCw8TCw9TCxMLExMLCx8LEwsvGwsjFwsTCxcPFwsTCxcLEwsfFwsTCxMPFwsTCxcfCxMLHwsTCwsTCwsTCxcfCxMLHy8XDwcFWVkJBV1JATMPBx8jBw8LCysfJxcfIycHHyMrIxsbTy8rAh4aHwMDHzYCHW+eAh26Ar4CAwMfMgIdbzoCFRkZHY4CkgIdTgJSAgMDrgL6AwVLBU0dQgJGAgVPBVEFUx1T8gIDB30lf/YCgYMRDQAdcgJ2Ag0xHdoC3gIdU+oCHXkKAyMLBzEBAAAAAAAAAIAAAAAAAAAAgAAAAAAAAAAjCwUhgAAAAAAAAACAAAAAAAAAAA0vAwMf5gIDAxID/gMdcVoCAwM/YgIFVR1xZgIdd4YCHRoDHgNhZmZpbmVfbWFwPChkMCwgZDEsIGQyKSAtPiAoZDAsIGQxLCBkMik+AGFmZmluZV9tYXA8KGQwLCBkMSkgLT4gKGQwLCBkMSk+AAVXAwM/VwMDP5oCHXmyAgVZAwMfBgMRCwEFWwVdBV8FYQVjBWUFZwVpBWsFbQVvBXEFcwMDH24CHXd+AgV1BXcDB30lf8YCgYMFeQV7BX0jAQEBBX8d+gL+AgMDHxYDBYEFgwWFBYcDAx9uAx2iA6YDBYkFiyMLAxEBAAAAAAAAAB2yA7YDHb4DwgMdygPOAx3eA+IDAwOnqQWNEQs9AxGtrxWxs7W3ubtXF72//8HDBY8BB/v5+Q0tBZEjCwcxIAAAAAAAAAABAAAAAAAAAAEAAAAAAAAABZMRCwkFlQWXBZkFmwEXxcnN0dXb3+Pn6+8DBRvHHTEJWQMFG8sdMQlbAwUbzx0xCV0DBRvTHTMJXwMFG9cd2QlhIwsFIQgAAAAAAAAAgAAAAAAAAAADBRvdHTMJYwMHS/cb4R0zCWUDB0v3G+UdMwlnAwdL9xvpHTMJaQMFG+0dMQlrAwUb8R0xCW0DBRU1F1kjdHB1Lm1lbW9yeV9zcGFjZTx2bWVtPgAjdHB1LnBpcGVsaW5lX21vZGU8c3luY2hyb25vdXM+ACN0cHUuZGltZW5zaW9uX3NlbWFudGljczxhcmJpdHJhcnk+ACN0cHUuZGltZW5zaW9uX3NlbWFudGljczxwYXJhbGxlbD4AI3RwdS5tZW1vcnlfc3BhY2U8c21lbT4AI3RwdS5jb3JlX3R5cGU8dGM+ACN0cHUuZG90X2RpbWVuc2lvbl9udW1iZXJzPFsxXSwgWzFdLCBbMF0sIFswXSwgWzAsIDAsIDEsIDBdLCBbXSwgW10+ACN0cHUuZG90X2RpbWVuc2lvbl9udW1iZXJzPFsxXSwgWzBdLCBbMF0sIFsxXSwgWzAsIDAsIDEsIDFdLCBbXSwgW10+AAMFFTUXWwMFFTUXXQMFFSkXXwMFFSkXYQMFFSkXYwMFFSkXZQMFFSkXZwMFFSkXaQMFFTUXawMFFTUXbREBAREDAR0+AhkFnQWfLUoCCe4lEzImBwWhBaMdVgINBaUdXgIZBacRCxEdagINBakRAREFqx16AhkFrR2CAhkFrx2KAg0FsQWzHZYCGQW1EQsFHaICDQW3AwMfqgITCRAAAOAPBbkdtgINBbsFvR3CAg0FvyMBCSEBAAAAAQAAAAMAAAAAAAAAHc4C0gIFwR3WAg0FwwXFHeICDQXHEwkBHe4CGQXJHYUNIwEJIQEAAAABAAAAAgAAAAAAAAAFyx0CAxkFzREBAgQdDgMZBc8F0SUFCQAAAAAF0x0iAw0F1QMJiwICjSWPJZElAwMuAzIDBdcjAQMJAQAAAB06Az4DBdkdQgMNBdsDAz9KAxELFR1SA1YDBd0dWgMNBd8dYgNmAwXhHWoDDQXjEwmQzMzMPx12A3oDBeUdfgMZBecdU4YDHYUZHY4DkgMF6R2WAwEF6wMDH54DJRcJAACA/wXtHaoDDQXvAwWXAgSZmwXxHboDDQXzBfUdxgMNBfcF+R3SAw0F+wMDH9oDJRcJAAAAAAX9HeYDDQX/AwWXBgSZmwMJiwYCjSWPJZElAwMf9gMRAQUjYXJpdGguZmFzdG1hdGg8bm9uZT4AI2FyaXRoLm92ZXJmbG93PG5vbmU+ACN2ZWN0b3Iua2luZDxtYXhpbXVtZj4AI3ZlY3Rvci5raW5kPGFkZD4AAQICAycFAgQCBAkX/QcFBQUPRwsBAgQBCQFBF/UHBQIEAgQfRycFAgQCBAEX9QUCBAIECUknAwIECScFAgQCBB8nBwUCBAIEHxf1BQIEAgQBSQcnBQIEAgQNF/UFIQIEAUkX9QcFAgQCBAlHJwcFAgQCBAknBQUCBAEnBQIEBQkFIQEBAQcHERERHSMdFRUVESUBBQsBAQEHBwcBAQEFCwEBAQcHBQEBBN4iBQERAaUHAwExCxEBqwcDU3chAQEBAQEBBwEHAREBEQERAR0BIwEdARUBFQEVAREBJQEDA3UJAwERB0NNAw0FBSEbBhEDAQMjAwMPCQMBEQcPTwMNBSUnHRQPAykJAx9NAwMtNwMJDwYhAwUDUwMDBwMDAwMDBwMDAwUGBwMFBxtXWRMFByMJVRtXWQMDLZMDCQ8GIQMFA10DAwcDAwMDAwcDAwMFBgcDBQcXYWMTBQcjCV8XYWMDAy03AwkPBiEDBQNnAwMHAwMDAwMHAwMDBQYHAwUHGWttEwUHIwlpGWttGQAPAwEFGQAPAwMLAwMDBwYLAwMDAwcGCwMDAwUVBgsDDwkJKy0vFwYRAwEDMQMDOwkDAREHQT0DDQUzNQMDCwMDAwcGCwMDAwMHBgsDAwMFFQYLAw8JBzk7PRcGEQMBAz8bBhEDAQM3AwMPCQMBEQcPTwMNBUNFHRQPA0cJA7PqAgMDhwkDAQMDL1UDASMHLzkDAQVTVQMDBQMDAwMDBQMDAwUGBQMFBxdZWwkGBQMFA10DAwUDAwMDAwUDAwMFBgUDBQcZYWMJBgUDBQNlAwMFAwMDAwMFAwMDAwMFAwMDBQYFAxsJC2lrbQkGBQMZA28DAwUDAwMHBgUDAwNXAwMFAwMDBQYFAxsJDXN1dwkGBQMZA3kDA0WJAwUnB0UmAwMFB3F7fQMDL1UDASMHLzkDAQVBgQMDL1UDASMHLzkDAQVThSkHKzkDAQWDhzcDNgMqAwMTDwYrAxMDiSkHKzkDEwWNiwMDBQMDAwMDBQMDAwUGBQMTBxWRkwkGBQMTA5URB04DRgMDIQWXjwMDBQMDAwcGBQMDA1cFBgUDKQcTm50JBgUDKQOfAwMFAwMDAwMFAwMDBQYFAxMHEaOlCQYFAxMDpw8GQwMTA6ERB0NNAyEFqas5Bl4DAyEFma0DA3IDkwMJDwaCAwMFA7E7BooDAwUHr3+zAwOVmgMDFysHla4DAxcFtbcJBiEDKwO5DwadAwUDuz0HnRMDBQVfvS0HnxMDBQW1vy8HoRMDBQPBAwOj1gMDFysHo+oDAxcFw8UJBiEDKwPHDwYhAwUDyS0HnxMDBQVfvy8HoRMDBQPNHwdREwMFBc9nIQcrEwMFBcvRAwMHAwMDAwMHAwMDBQYHAwUHF9XXEwUHIwm/F9XXAwMHAwMDAwMHAwMDBQYHAwUHGdvdEwUHIwnTGdvdAwMFAwMDBwYFAwMDVwMDBQMDAwUGBQMbCQ/h4+UJBgUDGQPnPwYRAwUD6QMDRYkDBScHRe4DAwUHw+vtAwMFAwMDAwMFAwMDBQYFAwUHG/HzCQYFAwUD9R8HURMDBQXP9yEHKxMDBQX57wMDBwMDAwMDBwMDAwUGBwMFBxv9/xMFByMJ+xv9/wMDh/IDAwEZAA8DAQUZAA8DA3UJAwERB0NNAw0FBUkbBhEDAQNLAwMPCQMBEQcPTwMNBU1PHRQPA1EJA13RAwMFAwMDAwMFAwMDBQYFAwUHGVNVCQYFAwUDVwMDJ6YCAwkPBicDBQNbMQcnEwMFBV1ZAwMFAwMDAwMFAwMDBQYFAwUHG2FjCQYFAwUDZR8HURMDBQVnXzMGEQMZA2kDAwcDAwMDAwcDAwMDAwcDAwMFBgcDGwkdbW9xCQYHAxkDcwkGBwMbA2sTBQd7C3cdbW9xNQfKAhMDBQNZAwMFAwMDAwMFAwMDBQYFAwUHF3t9CQYFAwUDfyEHKxMDBQV5gQMDBwMDAwMDBwMDAwMDBwMDAwUGBwMnCR+Fh4kJBgcDBQOLCQYHAycDgxMFB3sLjx+Fh4kDAy03AwkPBiEDBQORAwMHAwMDAwMHAwMDBQYHAwUHF5WXEwUHIwmTF5WXAwMtNwMJDwYhAwUDmwMDBwMDAwMDBwMDAwUGBwMFBxmfoRMFByMJnRmfoQMDLTcDCQ8GIQMFA6UDAwcDAwMDAwcDAwMFBgcDBQcbqasTBQcjCacbqasZAA8DAQUZAA8NAAELEQHzBwMPDwsBAQEBAQEHAQcBAwMBCQMBAwMBCQMBDQQBBwEDCwsRAQoCBwMrRwsBAQEBAQEHAQcBAwMLAwMDBwYLAwMDAwcGCwMDAwUVBgsDDwkJCw0PFwYRAwEDEQMDOwkDAREHQT0DDQUTFQMDCwMDAwcGCwMDAwMHBgsDAwMFFQYLAw8JBxkbHRcGEQMBAx8DAydzAwElBicDAQUBIwMDAQkDAQMDAQkDAQ0EAQclIScLEQEOAgcDK0cLAQEBAQEBBwEHAQMDCwMDAwcGCwMDAwMHBgsDAwMFFQYLAw8JCQsNDxcGEQMBAxEDAzsJAwERB0E9Aw0FExUDAwsDAwMHBgsDAwMDBwYLAwMDBRUGCwMPCQcZGx0XBhEDAQMfAwMncwMBJQYnAwEFASMDAwEJAwEDAwEJAwENBAEHJSEnCxEBEgIHAw8PCwEBAQEBAQcBBwEDAwEJAwEDAwEJAwENBAEFAwsLEQEWAgcDJz8LAQEBAQEBBwEHAQMDCwMDAwcGCwMDAwMHBgsDAwMFFQYLAw8JCQsNDxcGEQMBAxEDAzsJAwERB0E9Aw0FExUDAwsDAwMHBgsDAwMDBwYLAwMDBRUGCwMPCQcZGx0XBhEDAQMfAwMBCQMBAwMBCQMBDQQBBSMhCxEBGgIHAw8PCwEBAQEBAQcBBwEDAwEJAwEDAwEJAwENBAEFAwsLEQEeAgcDERMLAQEBAQEBBwEHAQMDAQkDAQMDAQkDAQMDAQkDAQ0EAQULDQsRASICBwMREwsBAQEBAQEHAQcBAwMBCQMBAwMBCQMBAwMBCQMBDQQBBQsNCxEBJgIHAxETCwEBAQEBAQcBBwEDAwEJAwEDAwEJAwEDAwEJAwENBAEFCw0LEQEqAgcDDw8LAQEBAQEBBwEHAQMDAQkDAQMDAQkDAQ0EAQcBAwsLEQEuAgcDDw8LAQEBAQEBBwEHAQMDAQkDAQMDAQkDAQ0EAQcBAwsGAwEFAQDmGgICLRkRCxELEQstGSUVCQsRCw0JFQ0XMRsdCQsNIxELEQsVDRETEQsNDQcJCw0HVS0SAiUJHR1HIyEjKS0fCx0nHSVFESkJCwkJCxsZGRkZGRkZGRkZJR0VDR0lEx0XHxsXExMbFxMXLxcXGRcXFw8ZFRkZIxcjGRUlIxkfDw8NCR0RYnVpbHRpbgBzdGFibGVfbW9zYWljAHRwdQBhcml0aAB2ZWN0b3IAbW9kdWxlAGFyaXRoLmNvbnN0YW50AHZlY3Rvci5sb2FkAGFyaXRoLmluZGV4X2Nhc3QAdmVjdG9yLnNoYXBlX2Nhc3QAZnVuYy5mdW5jAGZ1bmMucmV0dXJuAHZlY3Rvci5icm9hZGNhc3QAYXJpdGguY21waQB0cHUudmVjdG9yX3N0b3JlAG1lbXJlZi5sb2FkAGFyaXRoLmV4dHNpAHNjZi55aWVsZABhcml0aC5leHR1aQBzY2YuaWYAYXJpdGgubXVsZgBhcml0aC5hZGRmAGFyaXRoLm11bGkAYXJpdGguZGl2c2kAdHB1Lm1hdG11bABhcml0aC5hZGRpAHZlY3Rvci5tdWx0aV9yZWR1Y3Rpb24AYXJpdGguc3ViZgBtYXRoLmV4cABhcml0aC5kaXZmAGFyaXRoLnRydW5jZgBtYXRoLmxvZwB0cHUuaW90YQBhcml0aC5hbmRpAGFyaXRoLnNlbGVjdABhcml0aC5tYXhpbXVtZgBhcml0aC5leHRmAGZ1bmN0aW9uX3R5cGUAc3ltX25hbWUAdHJhbnNmb3JtX2luZGljZXMAd2luZG93X2JvdW5kcwB2YWx1ZQBwcmVkaWNhdGUAcGlwZWxpbmVfbW9kZQBicm9hZGNhc3RfaW5fZGltOgB0cmFuc2Zvcm1fMAB0cmFuc2Zvcm1fMQB0cmFuc2Zvcm1fMgB0cmFuc2Zvcm1fMwB0cmFuc2Zvcm1fNAB0cmFuc2Zvcm1fNQB0cmFuc2Zvcm1fNgB0cmFuc2Zvcm1fNwB0cmFuc2Zvcm1fOAB0cmFuc2Zvcm1fOQB0cmFuc2Zvcm1fMTAAZ2V0OgBndDoAZXE6AG11bDoAYWRkAG9wZXJhbmRTZWdtZW50U2l6ZXMAc3RyaWRlcwBicm9hZGNhc3RfaW5fZGltL2Jyb2FkY2FzdF9pbl9kaW0AZGltZW5zaW9uX251bWJlcnMAdHJhbnNwb3NlX2xocwB0cmFuc3Bvc2VfbGhzX2hpbnQAdHJhbnNwb3NlX3JocwBraW5kAHJlZHVjdGlvbl9kaW1zAHN0YWJsZV9tb3NhaWMudmVyc2lvbgBkaW1lbnNpb25fc2VtYW50aWNzAGl0ZXJhdGlvbl9ib3VuZHMAc2NhbGFyX3ByZWZldGNoAHNjcmF0Y2hfb3BlcmFuZHMAc3BsYXNoX21oYV9md2Rfc2VnbWVudGVkX3Jlc2lkdWFscwB0cHUuY29yZV90eXBlAHdpbmRvd19wYXJhbXMAZ2V0AF9zcGxhc2hfYXR0ZW50aW9uAC9fX3cvbWF4dGV4dC9tYXh0ZXh0Ly52ZW52L2xpYi9weXRob24zLjEyL3NpdGUtcGFja2FnZXMvamF4L2V4cGVyaW1lbnRhbC9wYWxsYXMvb3BzL3RwdS9zcGxhc2hfYXR0ZW50aW9uL3NwbGFzaF9hdHRlbnRpb25fa2VybmVsLnB5AGNvbnZlcnRfZWxlbWVudF90eXBlOgBjb252ZXJ0X2VsZW1lbnRfdHlwZS9jb252ZXJ0X2VsZW1lbnRfdHlwZQBndABndC9ndABkaXY6AGRpdgBlcQBlcS9lcQBjb25kOgBjb25kAGdldC9nZXQAZmFzdG1hdGgAbXVsL211bABzd2FwOgBzd2FwL3N3YXAAbG9nOgBsb2cvbG9nAGFkZDoAYWRkL2FkZABicm9hZGNhc3RfaW5fZGltAHNjYW46AHNjYW4AbXVsAG92ZXJmbG93RmxhZ3MAZG90X2dlbmVyYWw6AGRvdF9nZW5lcmFsL2RvdF9nZW5lcmFsAGRpbWVuc2lvbnMAaW90YToAaW90YS9pb3RhAGdlOgBnZS9nZQBhbmQ6AGFuZC9hbmQAaml0OgBqaXQAc2VsZWN0X246AHNlbGVjdF9uL3NlbGVjdF9uAHJlZHVjZV9tYXg6AHJlZHVjZV9tYXgvcmVkdWNlX21heABtYXg6AG1heC9tYXgAc3ViOgBzdWIvc3ViAGV4cDoAZXhwL2V4cAByZWR1Y2Vfc3VtOgByZWR1Y2Vfc3VtL3JlZHVjZV9zdW0A","needs_layout_passes":true,"allow_input_fusion":[],"serialization_format":"1","output_memory_colors":[],"output_memory_space_colors":[],"input_memory_space_colors":[]},"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %get-tuple-element.2190 = bf16[32,128,128]{2,1,0:T(8,128)(2,1)S(1)} get-tuple-element(%splash_mha_fwd_segmented_residuals.3), index=3, frontend_attributes={kernel_metadata={ +"xprof_metadata":"{\"block_q\": 128, \"block_kv\": 128, \"block_kv_compute\": 128, \"block_q_dkv\": 128, \"block_kv_dkv\": 128, \"block_kv_dkv_compute\": 128, \"block_q_dq\": 128, \"block_kv_dq\": 128, \"use_fused_bwd_kernel\": false, \"q_layout\": 1, \"k_layout\": 1, \"v_layout\": 1}" +}}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/vmap(static_sparsity_splash)/jit(_splash_attention)/splash_mha_fwd_segmented_residuals/splash_mha_fwd_segmented_residuals/pallas_call" stack_frame_id=0} + %slice-done.24 = s8[4194304]{0:T(1024)(128)(4,1)S(1)} async-done(%slice-start.24) + %slice-done.25 = s8[4194304]{0:T(1024)(128)(4,1)S(1)} async-done(%slice-start.25) + %slice-done.26 = s8[4194304]{0:T(1024)(128)(4,1)S(1)} async-done(%slice-start.26) + %slice-done.27 = s8[4194304]{0:T(1024)(128)(4,1)S(1)} async-done(%slice-start.27) + %slice-done.28 = s8[4194304]{0:T(1024)(128)(4,1)S(1)} async-done(%slice-start.28) + %slice-done.29 = s8[4194304]{0:T(1024)(128)(4,1)S(1)} async-done(%slice-start.29) + %slice-done.30 = s8[4194304]{0:T(1024)(128)(4,1)S(1)} async-done(%slice-start.30) + %slice-done.31 = s8[4194304]{0:T(1024)(128)(4,1)S(1)} async-done(%slice-start.31) + %custom-call.24 = bf16[1,32,128,4096]{3,2,1,0:T(8,128)(2,1)S(1)} custom-call(%slice-done.24, %slice-done.25, %slice-done.26, %slice-done.27, %slice-done.28, /*index=5*/%slice-done.29, %slice-done.30, %slice-done.31), custom_call_target="ConcatBitcast", backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %bitcast.913 = bf16[1,32,128,4096]{3,2,1,0:T(8,128)(2,1)S(1)} bitcast(%custom-call.24) + %fusion.505 = (f32[128]{0:T(128)}, bf16[1,128,4096]{2,1,0:T(8,128)(2,1)}) fusion(%get-tuple-element.2249, %bitcast.913, %get-tuple-element.2190), kind=kOutput, calls=%fused_computation.42.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %get-tuple-element.2191 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} get-tuple-element(%fusion.505), index=1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/reduce_sum" stack_frame_id=0} + %get-tuple-element.2333 = bf16[4,14336,1024]{2,1,0:T(8,128)(2,1)} get-tuple-element(%wide.param.1), index=31 + %dynamic-slice_convert_fusion.39 = bf16[1,14336,1024]{2,1,0:T(8,128)(2,1)S(1)} fusion(%get-tuple-element.2333, %copy.283), kind=kLoop, calls=%fused_computation.11.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %all-gather.103 = bf16[1,14336,4096]{2,1,0:T(8,128)(2,1)} all-gather(%dynamic-slice_convert_fusion.39), channel_id=106, replica_groups=mesh['axis_0'=1,'axis_1'=4] {'axis_1'}, dimensions={2}, use_global_device_ids=true, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"barrier_config":{"barrier_type":"CUSTOM","id":"0"},"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %get-tuple-element.2332 = bf16[4,1024,14336]{2,1,0:T(8,128)(2,1)} get-tuple-element(%wide.param.1), index=30 + %dynamic-slice_convert_fusion.40 = bf16[1,1024,14336]{2,1,0:T(8,128)(2,1)S(1)} fusion(%get-tuple-element.2332, %copy.283), kind=kLoop, calls=%fused_computation.12.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %all-gather.104 = bf16[1,4096,14336]{2,1,0:T(8,128)(2,1)} all-gather(%dynamic-slice_convert_fusion.40), channel_id=105, replica_groups=mesh['axis_0'=4,'axis_1'=1] {'axis_0'}, dimensions={1}, use_global_device_ids=true, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"barrier_config":{"barrier_type":"CUSTOM","id":"0"},"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %get-tuple-element.2331 = bf16[4,1024,14336]{2,1,0:T(8,128)(2,1)} get-tuple-element(%wide.param.1), index=29 + %dynamic-slice_convert_fusion.41 = bf16[1,1024,14336]{2,1,0:T(8,128)(2,1)S(1)} fusion(%get-tuple-element.2331, %copy.283), kind=kLoop, calls=%fused_computation.13.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %get-tuple-element.2192 = f32[128]{0:T(128)} get-tuple-element(%fusion.505), index=0, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/reduce_sum" stack_frame_id=0} + %copy-start.127 = (f32[128]{0:T(128)S(1)}, f32[128]{0:T(128)}, u32[]{:S(2)}) copy-start(%get-tuple-element.2192) + %copy-start.3 = (bf16[1,128,4096]{2,1,0:T(8,128)(2,1)S(1)}, bf16[1,128,4096]{2,1,0:T(8,128)(2,1)}, u32[]{:S(2)}) copy-start(%get-tuple-element.2191) + %get-tuple-element.2193 = bf16[4096]{0:T(1024)(128)(2,1)} get-tuple-element(%convert_reduce_fusion.9), index=1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %copy-start.50 = (bf16[4096]{0:T(1024)(128)(2,1)S(1)}, bf16[4096]{0:T(1024)(128)(2,1)}, u32[]{:S(2)}) copy-start(%get-tuple-element.2193) + %all-gather.105 = bf16[1,4096,14336]{2,1,0:T(8,128)(2,1)} all-gather(%dynamic-slice_convert_fusion.41), channel_id=104, replica_groups=mesh['axis_0'=4,'axis_1'=1] {'axis_0'}, dimensions={1}, use_global_device_ids=true, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"barrier_config":{"barrier_type":"CUSTOM","id":"0"},"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %copy-done.127 = f32[128]{0:T(128)S(1)} copy-done(%copy-start.127) + %add_rsqrt_fusion.5 = f32[128]{0:T(128)S(1)} fusion(%copy-done.127), kind=kLoop, calls=%fused_computation.68.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/rsqrt" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %copy-done.3 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)S(1)} copy-done(%copy-start.3) + %copy-done.50 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} copy-done(%copy-start.50) + %get-tuple-element.2334 = u32[4]{0:T(128)} get-tuple-element(%wide.param.1), index=32, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_update_index_in_dim)/dynamic_update_slice" stack_frame_id=0} + %copy-start.143 = (u32[4]{0:T(128)S(1)}, u32[4]{0:T(128)}, u32[]{:S(2)}) copy-start(%get-tuple-element.2334) + %fusion.506 = bf16[1,128,14336]{2,1,0:T(8,128)(2,1)S(1)} fusion(%all-gather.105, %copy-done.3, %add_rsqrt_fusion.5, %copy-done.50), kind=kOutput, calls=%fused_computation.24.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %get-tuple-element.2335 = u32[4,4]{1,0:T(4,128)} get-tuple-element(%wide.param.1), index=33, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_update_index_in_dim)/dynamic_update_slice" stack_frame_id=0} + %copy-start.128 = (u32[4,4]{1,0:T(4,128)S(1)}, u32[4,4]{1,0:T(4,128)}, u32[]{:S(2)}) copy-start(%get-tuple-element.2335) + %get-tuple-element.2336 = u32[4]{0:T(128)} get-tuple-element(%wide.param.1), index=34, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_update_index_in_dim)/dynamic_update_slice" stack_frame_id=0} + %copy-start.144 = (u32[4]{0:T(128)S(1)}, u32[4]{0:T(128)}, u32[]{:S(2)}) copy-start(%get-tuple-element.2336) + %get-tuple-element.2337 = u32[4,4]{1,0:T(4,128)} get-tuple-element(%wide.param.1), index=35, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_update_index_in_dim)/dynamic_update_slice" stack_frame_id=0} + %copy-start.129 = (u32[4,4]{1,0:T(4,128)S(1)}, u32[4,4]{1,0:T(4,128)}, u32[]{:S(2)}) copy-start(%get-tuple-element.2337) + %get-tuple-element.2338 = u32[4]{0:T(128)} get-tuple-element(%wide.param.1), index=36, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_update_index_in_dim)/dynamic_update_slice" stack_frame_id=0} + %copy-start.145 = (u32[4]{0:T(128)S(1)}, u32[4]{0:T(128)}, u32[]{:S(2)}) copy-start(%get-tuple-element.2338) + %get-tuple-element.2339 = u32[4,4]{1,0:T(4,128)} get-tuple-element(%wide.param.1), index=37, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_update_index_in_dim)/dynamic_update_slice" stack_frame_id=0} + %copy-start.130 = (u32[4,4]{1,0:T(4,128)S(1)}, u32[4,4]{1,0:T(4,128)}, u32[]{:S(2)}) copy-start(%get-tuple-element.2339) + %get-tuple-element.2340 = u32[4]{0:T(128)} get-tuple-element(%wide.param.1), index=38, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_update_index_in_dim)/dynamic_update_slice" stack_frame_id=0} + %copy-start.146 = (u32[4]{0:T(128)S(1)}, u32[4]{0:T(128)}, u32[]{:S(2)}) copy-start(%get-tuple-element.2340) + %get-tuple-element.2341 = u32[4,4]{1,0:T(4,128)} get-tuple-element(%wide.param.1), index=39, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_update_index_in_dim)/dynamic_update_slice" stack_frame_id=0} + %copy-start.131 = (u32[4,4]{1,0:T(4,128)S(1)}, u32[4,4]{1,0:T(4,128)}, u32[]{:S(2)}) copy-start(%get-tuple-element.2341) + %get-tuple-element.2342 = u32[4]{0:T(128)} get-tuple-element(%wide.param.1), index=40, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_update_index_in_dim)/dynamic_update_slice" stack_frame_id=0} + %copy-start.147 = (u32[4]{0:T(128)S(1)}, u32[4]{0:T(128)}, u32[]{:S(2)}) copy-start(%get-tuple-element.2342) + %get-tuple-element.2343 = u32[4,4]{1,0:T(4,128)} get-tuple-element(%wide.param.1), index=41, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_update_index_in_dim)/dynamic_update_slice" stack_frame_id=0} + %copy-start.132 = (u32[4,4]{1,0:T(4,128)S(1)}, u32[4,4]{1,0:T(4,128)}, u32[]{:S(2)}) copy-start(%get-tuple-element.2343) + %get-tuple-element.2344 = u32[4]{0:T(128)} get-tuple-element(%wide.param.1), index=42, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_update_index_in_dim)/dynamic_update_slice" stack_frame_id=0} + %copy-start.148 = (u32[4]{0:T(128)S(1)}, u32[4]{0:T(128)}, u32[]{:S(2)}) copy-start(%get-tuple-element.2344) + %get-tuple-element.2345 = u32[4,4]{1,0:T(4,128)} get-tuple-element(%wide.param.1), index=43, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_update_index_in_dim)/dynamic_update_slice" stack_frame_id=0} + %copy-start.133 = (u32[4,4]{1,0:T(4,128)S(1)}, u32[4,4]{1,0:T(4,128)}, u32[]{:S(2)}) copy-start(%get-tuple-element.2345) + %get-tuple-element.2346 = u32[4]{0:T(128)} get-tuple-element(%wide.param.1), index=44, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_update_index_in_dim)/dynamic_update_slice" stack_frame_id=0} + %copy-start.149 = (u32[4]{0:T(128)S(1)}, u32[4]{0:T(128)}, u32[]{:S(2)}) copy-start(%get-tuple-element.2346) + %get-tuple-element.2347 = u32[4,4]{1,0:T(4,128)} get-tuple-element(%wide.param.1), index=45, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_update_index_in_dim)/dynamic_update_slice" stack_frame_id=0} + %copy-start.134 = (u32[4,4]{1,0:T(4,128)S(1)}, u32[4,4]{1,0:T(4,128)}, u32[]{:S(2)}) copy-start(%get-tuple-element.2347) + %get-tuple-element.2348 = u32[4]{0:T(128)} get-tuple-element(%wide.param.1), index=46, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_update_index_in_dim)/dynamic_update_slice" stack_frame_id=0} + %copy-start.150 = (u32[4]{0:T(128)S(1)}, u32[4]{0:T(128)}, u32[]{:S(2)}) copy-start(%get-tuple-element.2348) + %get-tuple-element.2349 = u32[4,4]{1,0:T(4,128)} get-tuple-element(%wide.param.1), index=47, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_update_index_in_dim)/dynamic_update_slice" stack_frame_id=0} + %copy-start.135 = (u32[4,4]{1,0:T(4,128)S(1)}, u32[4,4]{1,0:T(4,128)}, u32[]{:S(2)}) copy-start(%get-tuple-element.2349) + %get-tuple-element.2350 = u32[4]{0:T(128)} get-tuple-element(%wide.param.1), index=48, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_update_index_in_dim)/dynamic_update_slice" stack_frame_id=0} + %copy-start.151 = (u32[4]{0:T(128)S(1)}, u32[4]{0:T(128)}, u32[]{:S(2)}) copy-start(%get-tuple-element.2350) + %get-tuple-element.2351 = u32[4,4]{1,0:T(4,128)} get-tuple-element(%wide.param.1), index=49, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_update_index_in_dim)/dynamic_update_slice" stack_frame_id=0} + %copy-start.136 = (u32[4,4]{1,0:T(4,128)S(1)}, u32[4,4]{1,0:T(4,128)}, u32[]{:S(2)}) copy-start(%get-tuple-element.2351) + %get-tuple-element.2767 = bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)} get-tuple-element(%wide.param.1), index=54 + %copy-start.75 = (bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)S(1)}, bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)}, u32[]{:S(2)}) copy-start(%get-tuple-element.2767) + %get-tuple-element.2823 = bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)} get-tuple-element(%wide.param.1), index=55 + %copy-start.77 = (bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)S(1)}, bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)}, u32[]{:S(2)}) copy-start(%get-tuple-element.2823) + %get-tuple-element.2880 = s32[1,128]{1,0:T(1,128)} get-tuple-element(%wide.param.1), index=56 + %copy-start.106 = (s32[1,128]{1,0:T(1,128)S(1)}, s32[1,128]{1,0:T(1,128)}, u32[]{:S(2)}) copy-start(%get-tuple-element.2880) + %fusion.507 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)S(1)} fusion(%copy-done.3, %all-gather.103, %all-gather.104, %fusion.506, %add_rsqrt_fusion.5, /*index=5*/%copy-done.50), kind=kOutput, calls=%fused_computation.38.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %copy-done.143 = u32[4]{0:T(128)S(1)} copy-done(%copy-start.143) + %copy.306 = u32[4]{0:T(128)} copy(%copy-done.143), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %copy-done.128 = u32[4,4]{1,0:T(4,128)S(1)} copy-done(%copy-start.128) + %copy.307 = u32[4,4]{1,0:T(4,128)} copy(%copy-done.128), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %copy-done.144 = u32[4]{0:T(128)S(1)} copy-done(%copy-start.144) + %copy.308 = u32[4]{0:T(128)} copy(%copy-done.144), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %copy-done.129 = u32[4,4]{1,0:T(4,128)S(1)} copy-done(%copy-start.129) + %copy.309 = u32[4,4]{1,0:T(4,128)} copy(%copy-done.129), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %copy-done.145 = u32[4]{0:T(128)S(1)} copy-done(%copy-start.145) + %copy.310 = u32[4]{0:T(128)} copy(%copy-done.145), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %copy-done.130 = u32[4,4]{1,0:T(4,128)S(1)} copy-done(%copy-start.130) + %copy.311 = u32[4,4]{1,0:T(4,128)} copy(%copy-done.130), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %copy-done.146 = u32[4]{0:T(128)S(1)} copy-done(%copy-start.146) + %copy.312 = u32[4]{0:T(128)} copy(%copy-done.146), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %copy-done.131 = u32[4,4]{1,0:T(4,128)S(1)} copy-done(%copy-start.131) + %copy.313 = u32[4,4]{1,0:T(4,128)} copy(%copy-done.131), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %copy-done.147 = u32[4]{0:T(128)S(1)} copy-done(%copy-start.147) + %copy.314 = u32[4]{0:T(128)} copy(%copy-done.147), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %copy-done.132 = u32[4,4]{1,0:T(4,128)S(1)} copy-done(%copy-start.132) + %copy.315 = u32[4,4]{1,0:T(4,128)} copy(%copy-done.132), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %copy-done.148 = u32[4]{0:T(128)S(1)} copy-done(%copy-start.148) + %copy.316 = u32[4]{0:T(128)} copy(%copy-done.148), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %copy-done.133 = u32[4,4]{1,0:T(4,128)S(1)} copy-done(%copy-start.133) + %copy.317 = u32[4,4]{1,0:T(4,128)} copy(%copy-done.133), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %copy-done.149 = u32[4]{0:T(128)S(1)} copy-done(%copy-start.149) + %copy.318 = u32[4]{0:T(128)} copy(%copy-done.149), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %copy-done.134 = u32[4,4]{1,0:T(4,128)S(1)} copy-done(%copy-start.134) + %copy.319 = u32[4,4]{1,0:T(4,128)} copy(%copy-done.134), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %copy-done.150 = u32[4]{0:T(128)S(1)} copy-done(%copy-start.150) + %copy.320 = u32[4]{0:T(128)} copy(%copy-done.150), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %copy-done.135 = u32[4,4]{1,0:T(4,128)S(1)} copy-done(%copy-start.135) + %copy.321 = u32[4,4]{1,0:T(4,128)} copy(%copy-done.135), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %copy-done.151 = u32[4]{0:T(128)S(1)} copy-done(%copy-start.151) + %copy.322 = u32[4]{0:T(128)} copy(%copy-done.151), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %copy-done.136 = u32[4,4]{1,0:T(4,128)S(1)} copy-done(%copy-start.136) + %copy.323 = u32[4,4]{1,0:T(4,128)} copy(%copy-done.136), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %get-tuple-element.2353 = s32[]{:T(128)} get-tuple-element(%wide.param.1), index=51 + %copy-done.75 = bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)S(1)} copy-done(%copy-start.75) + %copy-done.77 = bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)S(1)} copy-done(%copy-start.77) + %copy-done.106 = s32[1,128]{1,0:T(1,128)S(1)} copy-done(%copy-start.106) + %bitcast.943 = s32[128]{0:T(128)S(1)} bitcast(%copy-done.106) + ROOT %tuple.241 = (s32[]{:T(128)}, bf16[1,128,4096]{2,1,0:T(8,128)(2,1)S(1)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, /*index=5*/u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, /*index=10*/u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, /*index=15*/u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, /*index=20*/bf16[4,1,128,4096]{3,2,1,0:T(8,128)(2,1)}, bf16[4,4096]{1,0:T(4,128)(2,1)}, bf16[4,4096]{1,0:T(4,128)(2,1)}, bf16[4,1024,32,128]{1,3,2,0:T(8,128)(2,1)}, bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)S(1)}, /*index=25*/bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)S(1)}, bf16[4,1024,8,128]{1,3,2,0:T(8,128)(2,1)}, bf16[4,1024,8,128]{3,1,2,0:T(8,128)(2,1)}, bf16[4,32,128,1024]{3,2,1,0:T(8,128)(2,1)}, bf16[4,1024,14336]{2,1,0:T(8,128)(2,1)}, /*index=30*/bf16[4,1024,14336]{2,1,0:T(8,128)(2,1)}, bf16[4,14336,1024]{2,1,0:T(8,128)(2,1)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, /*index=35*/u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, /*index=40*/u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, /*index=45*/u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, /*index=50*/s32[128]{0:T(128)S(1)}, s32[]{:T(128)}, s8[1,1,1]{2,1,0:T(4,128)(4,1)}, s8[1,1,1]{2,1,0:T(4,128)(4,1)}, bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)}, /*index=55*/bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)}, s32[1,128]{1,0:T(1,128)}) tuple(%add.904, %fusion.507, %copy.306, %copy.307, %copy.308, /*index=5*/%copy.309, %copy.310, %copy.311, %copy.312, %copy.313, /*index=10*/%copy.314, %copy.315, %copy.316, %copy.317, %copy.318, /*index=15*/%copy.319, %copy.320, %copy.321, %copy.322, %copy.323, /*index=20*/%bitcast_dynamic-update-slice_fusion.5, %get-tuple-element.2323, %get-tuple-element.2324, %get-tuple-element.2325, %copy-done.75, /*index=25*/%copy-done.77, %get-tuple-element.2328, %get-tuple-element.2329, %get-tuple-element.2330, %get-tuple-element.2331, /*index=30*/%get-tuple-element.2332, %get-tuple-element.2333, %get-tuple-element.2334, %get-tuple-element.2335, %get-tuple-element.2336, /*index=35*/%get-tuple-element.2337, %get-tuple-element.2338, %get-tuple-element.2339, %get-tuple-element.2340, %get-tuple-element.2341, /*index=40*/%get-tuple-element.2342, %get-tuple-element.2343, %get-tuple-element.2344, %get-tuple-element.2345, %get-tuple-element.2346, /*index=45*/%get-tuple-element.2347, %get-tuple-element.2348, %get-tuple-element.2349, %get-tuple-element.2350, %get-tuple-element.2351, /*index=50*/%bitcast.943, %get-tuple-element.2353, %get-tuple-element.2354, %get-tuple-element.2355, %get-tuple-element.2767, /*index=55*/%get-tuple-element.2823, %get-tuple-element.2880) +} + +%wide.region_4.7_spmd.clone.clone.clone (wide.param.116: (s32[], bf16[1,128,4096], u32[4], u32[4,4], u32[4], /*index=5*/u32[4,4], u32[4], u32[4,4], u32[4], u32[4,4], /*index=10*/u32[4], u32[4,4], u32[4], u32[4,4], u32[4], /*index=15*/u32[4,4], u32[4], u32[4,4], u32[4], u32[4,4], /*index=20*/bf16[4,1,128,4096], bf16[4,4096], bf16[4,4096], bf16[4,1024,32,128], bf16[1,128,1,128], /*index=25*/bf16[1,128,1,128], bf16[4,1024,8,128], bf16[4,1024,8,128], bf16[4,32,128,1024], bf16[4,1024,14336], /*index=30*/bf16[4,1024,14336], bf16[4,14336,1024], u32[4], u32[4,4], u32[4], /*index=35*/u32[4,4], u32[4], u32[4,4], u32[4], u32[4,4], /*index=40*/u32[4], u32[4,4], u32[4], u32[4,4], u32[4], /*index=45*/u32[4,4], u32[4], u32[4,4], u32[4], u32[4,4], /*index=50*/s32[128], s32[], s8[1,1,1], s8[1,1,1], bf16[1,128,1,128], /*index=55*/bf16[1,128,1,128], s32[1,128])) -> pred[] { + %constant.358.clone.4 = s32[]{:T(128)} constant(4) + %wide.param.116 = (s32[]{:T(128)}, bf16[1,128,4096]{2,1,0:T(8,128)(2,1)S(1)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, /*index=5*/u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, /*index=10*/u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, /*index=15*/u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, /*index=20*/bf16[4,1,128,4096]{3,2,1,0:T(8,128)(2,1)}, bf16[4,4096]{1,0:T(4,128)(2,1)}, bf16[4,4096]{1,0:T(4,128)(2,1)}, bf16[4,1024,32,128]{1,3,2,0:T(8,128)(2,1)}, bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)S(1)}, /*index=25*/bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)S(1)}, bf16[4,1024,8,128]{1,3,2,0:T(8,128)(2,1)}, bf16[4,1024,8,128]{3,1,2,0:T(8,128)(2,1)}, bf16[4,32,128,1024]{3,2,1,0:T(8,128)(2,1)}, bf16[4,1024,14336]{2,1,0:T(8,128)(2,1)}, /*index=30*/bf16[4,1024,14336]{2,1,0:T(8,128)(2,1)}, bf16[4,14336,1024]{2,1,0:T(8,128)(2,1)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, /*index=35*/u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, /*index=40*/u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, /*index=45*/u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, /*index=50*/s32[128]{0:T(128)S(1)}, s32[]{:T(128)}, s8[1,1,1]{2,1,0:T(4,128)(4,1)}, s8[1,1,1]{2,1,0:T(4,128)(4,1)}, bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)}, /*index=55*/bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)}, s32[1,128]{1,0:T(1,128)}) parameter(0) + %get-tuple-element.1896 = s32[]{:T(128)} get-tuple-element(%wide.param.116), index=0 + ROOT %lt.34 = pred[]{:T(512)} compare(%get-tuple-element.1896, %constant.358.clone.4), direction=LT, metadata={op_name="jit(train_step)/jvp()/while/cond/lt" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } %region_5.8 (reduce_sum.87: f32[], reduce_sum.88: f32[]) -> f32[] { - %reduce_sum.87 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/reduce_sum"} - %reduce_sum.88 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/reduce_sum"} - ROOT %reduce_sum.92 = f32[]{:T(128)} add(%reduce_sum.87, %reduce_sum.88), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.312 (param_0.1131: bf16[4,128,4096]) -> f32[4,128] { - %param_0.1131 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) - %convert_element_type.975 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_0.1131), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %square.192 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.975, %convert_element_type.975), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/square" stack_frame_id=0} - %constant.1019 = f32[]{:T(128)} constant(0) - ROOT %reduce.147 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%square.192, %constant.1019), dimensions={2}, to_apply=%region_5.8, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/reduce_sum" stack_frame_id=0} -} - -%region_10.13 (reduce_sum.102: f32[], reduce_sum.106: f32[]) -> f32[] { - %reduce_sum.102 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/reduce_sum"} - %reduce_sum.106 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/reduce_sum"} - ROOT %reduce_sum.107 = f32[]{:T(128)} add(%reduce_sum.102, %reduce_sum.106), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.314 (param_0.1126: bf16[4,128,4096], param_1.1285: bf16[4,128,4096], param_2.1113: bf16[4096]) -> f32[4,128] { - %param_0.1126 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) - %convert_element_type.982 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_0.1126), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %param_1.1285 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %param_2.1113 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(2) - %dot_general.374 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1113), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %dot_general.364 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%param_1.1285, %dot_general.374), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %convert_element_type.981 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%dot_general.364), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %mul.1419 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.982, %convert_element_type.981), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %constant.1013 = f32[]{:T(128)} constant(0) - ROOT %reduce.148 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%mul.1419, %constant.1013), dimensions={2}, to_apply=%region_10.13, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/reduce_sum" stack_frame_id=0} + %reduce_sum.87 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} + %reduce_sum.88 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} + ROOT %reduce_sum.92 = f32[]{:T(128)} add(%reduce_sum.87, %reduce_sum.88), metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%fused_computation.352 (param_0.1293: bf16[1,128,4096]) -> f32[128] { + %param_0.1293 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.901 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%param_0.1293), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} + %square.193 = f32[1,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.901, %convert_element_type.901), metadata={op_name="jit(train_step)/jvp()/square" stack_frame_id=0} + %constant.389.clone.28 = f32[]{:T(128)} constant(0) + ROOT %reduce.160 = f32[128]{0:T(128)S(1)} reduce(%square.193, %constant.389.clone.28), dimensions={0,2}, to_apply=%region_5.8, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} +} + +%fused_computation.397 (param_0.1158: f32[128]) -> f32[128] { + %param_0.1158 = f32[128]{0:T(128)S(1)} parameter(0) + %constant.398.clone.2 = f32[]{:T(128)} constant(0.000244140625) + %broadcast.632 = f32[128]{0:T(128)} broadcast(%constant.398.clone.2), dimensions={}, metadata={op_name="broadcast.283"} + %div.667 = f32[128]{0:T(128)} multiply(%param_0.1158, %broadcast.632), metadata={op_name="jit(train_step)/jvp()/div" stack_frame_id=0} + %constant.399.clone.2 = f32[]{:T(128)} constant(1e-05) + %broadcast.630 = f32[128]{0:T(128)} broadcast(%constant.399.clone.2), dimensions={}, metadata={op_name="jit(train_step)/jvp()/add" stack_frame_id=0} + %add.791 = f32[128]{0:T(128)} add(%div.667, %broadcast.630), metadata={op_name="jit(train_step)/jvp()/add" stack_frame_id=0} + %bitcast.568 = f32[1,128]{1,0:T(1,128)} bitcast(%add.791), metadata={op_name="jit(train_step)/jvp()/add" stack_frame_id=0} + %rsqrt.101 = f32[1,128]{1,0:T(1,128)} rsqrt(%bitcast.568), metadata={op_name="jit(train_step)/jvp()/rsqrt" stack_frame_id=0} + ROOT %bitcast.558 = f32[128]{0:T(128)} bitcast(%rsqrt.101), metadata={op_name="jit(train_step)/jvp()/rsqrt" stack_frame_id=0} +} + +%fused_computation.358.clone.clone (param_0.1260: bf16[4096], param_1.1403: f32[128], param_2.1079: bf16[1,128,4096]) -> bf16[128,4096] { + %param_0.1260 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(0) + %dot_general.406 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_0.1260), dimensions={2}, metadata={op_name="jit(train_step)/jvp(...k,k->...k)/dot_general" stack_frame_id=0} + %convert.142 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%dot_general.406) + %param_2.1079 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} parameter(2) + %convert_element_type.960 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%param_2.1079), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} + %param_1.1403 = f32[128]{0:T(128)S(1)} parameter(1) + %mul.1644 = f32[1,128,4096]{2,1,0:T(8,128)} broadcast(%param_1.1403), dimensions={1}, metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + %mul.1643 = f32[1,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.960, %mul.1644), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + %convert_element_type.959 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1643), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} + %convert.143 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%convert_element_type.959) + %dot_general.405 = f32[1,128,4096]{2,1,0:T(8,128)} multiply(%convert.142, %convert.143), metadata={op_name="jit(train_step)/jvp(...k,k->...k)/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","4096"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %convert.144 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} convert(%dot_general.405) + ROOT %bitcast.649 = bf16[128,4096]{1,0:T(8,128)(2,1)} bitcast(%convert.144) +} + +%bitcast_fusion.5 (bitcast_input.5: bf16[4096,128256]) -> bf16[4096,128256] { + %bitcast_input.5 = bf16[4096,128256]{1,0:T(8,128)(2,1)} parameter(0) + ROOT %bitcast.661 = bf16[4096,128256]{1,0:T(8,128)(2,1)} bitcast(%bitcast_input.5) } -%region_8.11 (dot_general.182: bf16[], dot_general.183: bf16[]) -> bf16[] { - %dot_general.182 = bf16[]{:T(256)} parameter(0), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general"} - %dot_general.183 = bf16[]{:T(256)} parameter(1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general"} - ROOT %add.168 = bf16[]{:T(256)} add(%dot_general.182, %dot_general.183), metadata={op_name="add.54"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.235.clone.clone (param_0.1095: f32[4096,128256]) -> bf16[4096,128256,1] { - %param_0.1095 = f32[4096,128256]{1,0:T(8,128)} parameter(0) - %convert_element_type.1033 = bf16[4096,128256]{1,0:T(8,128)(2,1)} convert(%param_0.1095), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} - ROOT %bitcast.449 = bf16[4096,128256,1]{1,0,2:T(8,128)(2,1)} bitcast(%convert_element_type.1033), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} -} - -%fused_computation.280.clone.1.clone.clone (param_0.1096: bf16[4,128,128256], param_1.1261: s32[4,128], param_2.1081: f32[4,128], param_3.782: f32[4,128], param_4.484: bf16[4,128], param_5.409: f32[4,128]) -> bf16[4,128,128256] { - %param_5.409 = f32[4,128]{1,0:T(4,128)S(1)} parameter(5) - %mul.1603 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_5.409), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - %param_3.782 = f32[4,128]{1,0:T(4,128)S(1)} parameter(3) - %mul.1602 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_3.782), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - %param_0.1096 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} parameter(0) - %convert_element_type.1036 = f32[4,128,128256]{2,1,0:T(8,128)} convert(%param_0.1096), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %param_4.484 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(4) - %sub.88 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_4.484), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %sub.87 = f32[4,128,128256]{2,1,0:T(8,128)} subtract(%convert_element_type.1036, %sub.88), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %exp.60 = f32[4,128,128256]{2,1,0:T(8,128)} exponential(%sub.87), metadata={op_name="jit(train_step)/jvp()/exp" stack_frame_id=0} - %mul.1601 = f32[4,128,128256]{2,1,0:T(8,128)} multiply(%mul.1602, %exp.60), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - %param_2.1081 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %div.819 = f32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_2.1081), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} - %div.818 = f32[4,128,128256]{2,1,0:T(8,128)} divide(%mul.1601, %div.819), metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} - %param_1.1261 = s32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %eq.43 = s32[4,128,128256]{2,1,0:T(8,128)} broadcast(%param_1.1261), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} - %eq.42 = s32[4,128,128256]{2,1,0:T(8,128)} iota(), iota_dimension=2, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} - %eq.41 = pred[4,128,128256]{2,1,0:T(8,128)(4,1)} compare(%eq.43, %eq.42), direction=EQ, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} - %convert_element_type.1035 = f32[4,128,128256]{2,1,0:T(8,128)} convert(%eq.41), metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/convert_element_type" stack_frame_id=0} - %sub.86 = f32[4,128,128256]{2,1,0:T(8,128)} subtract(%div.818, %convert_element_type.1035), metadata={op_name="jit(train_step)/transpose(jvp())/sub" stack_frame_id=0} - %mul.1600 = f32[4,128,128256]{2,1,0:T(8,128)} multiply(%mul.1603, %sub.86), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - ROOT %convert_element_type.1034 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} convert(%mul.1600), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} -} - -%fused_computation.315 (param_0.1094: f32[4,128], param_1.1260: bf16[4,128,4096], param_2.1082: f32[4096,128256], param_3.783: bf16[4,128,128256], param_4.485: s32[4,128], param_5.410: f32[4,128], param_6.284: f32[4,128], param_7.183: bf16[4,128], param_8.102: f32[4,128]) -> (bf16[4096], bf16[4,128,4096]) { - %param_3.783 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} parameter(3) - %param_4.485 = s32[4,128]{1,0:T(4,128)S(1)} parameter(4) - %param_5.410 = f32[4,128]{1,0:T(4,128)S(1)} parameter(5) - %param_6.284 = f32[4,128]{1,0:T(4,128)S(1)} parameter(6) - %param_7.183 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(7) - %param_8.102 = f32[4,128]{1,0:T(4,128)S(1)} parameter(8) - %multiply_convert_fusion.2.clone.1 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} fusion(%param_3.783, %param_4.485, %param_5.410, %param_6.284, %param_7.183, /*index=5*/%param_8.102), kind=kLoop, calls=%fused_computation.280.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %param_2.1082 = f32[4096,128256]{1,0:T(8,128)} parameter(2) - %fusion.219.clone.1 = bf16[4096,128256,1]{1,0,2:T(8,128)(2,1)} fusion(%param_2.1082), kind=kLoop, calls=%fused_computation.235.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} - %convolution.86.clone.1 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} convolution(%multiply_convert_fusion.2.clone.1, %fusion.219.clone.1), window={size=1}, dim_labels=0bf_oi0->0bf, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/dot_general" stack_frame_id=0} - %param_1.1260 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %convert_element_type.994 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_1.1260), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %param_0.1094 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.1434 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_0.1094), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %mul.1433 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.994, %mul.1434), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %convert_element_type.993 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1433), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %multiply.252 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%convolution.86.clone.1, %convert_element_type.993), metadata={op_name="multiply.206"} - %constant.874 = bf16[]{:T(256)} constant(0) - %reduce.149 = bf16[4096]{0:T(1024)(128)(2,1)} reduce(%multiply.252, %constant.874), dimensions={0,1}, to_apply=%region_8.11, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - ROOT %tuple.153 = (bf16[4096]{0:T(1024)(128)(2,1)}, bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)}) tuple(%reduce.149, %convolution.86.clone.1) -} - -%fused_computation.323 (param_0.904: f32[64], param_1.974: f32[4,128]) -> (bf16[4,128,1,64], bf16[4,128,1,64]) { - %param_1.974 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %div.621 = f32[4,128,1,64]{3,1,0,2:T(8,128)} broadcast(%param_1.974), dimensions={0,1}, metadata={op_name="jit(train_step)/layers/div" stack_frame_id=0} - %param_0.904 = f32[64]{0:T(128)S(1)} parameter(0) - %div.619 = f32[4,128,1,64]{3,1,0,2:T(8,128)} broadcast(%param_0.904), dimensions={3}, metadata={op_name="jit(train_step)/layers/div" stack_frame_id=0} - %div.618 = f32[4,128,1,64]{3,1,0,2:T(8,128)} divide(%div.621, %div.619), metadata={op_name="jit(train_step)/layers/div" stack_frame_id=0} - %sin.38 = f32[4,128,1,64]{3,1,0,2:T(8,128)} sine(%div.618), metadata={op_name="jit(train_step)/layers/sin" stack_frame_id=0} - %convert_element_type.1002 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} convert(%sin.38), metadata={op_name="jit(train_step)/layers/convert_element_type" stack_frame_id=0} - %cos.41.clone.1 = f32[4,128,1,64]{3,1,0,2:T(8,128)} cosine(%div.618), metadata={op_name="jit(train_step)/layers/cos" stack_frame_id=0} - %convert_element_type.1001.clone.1 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} convert(%cos.41.clone.1), metadata={op_name="jit(train_step)/layers/convert_element_type" stack_frame_id=0} - ROOT %tuple.150 = (bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)}, bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)}) tuple(%convert_element_type.1002, %convert_element_type.1001.clone.1) -} - -%fused_computation.324 (param_0.901: bf16[4,128,1,64]) -> bf16[4,128,1,128] { - %param_0.901 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} parameter(0) - %constant.866 = bf16[]{:T(256)} constant(-inf) - %pad.38 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.901, %constant.866), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} - %pad.37 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.901, %constant.866), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} - ROOT %maximum.34 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} maximum(%pad.38, %pad.37), metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} -} - -%fused_computation.325 (param_0.903: bf16[4,128,1,64]) -> bf16[4,128,1,128] { - %param_0.903 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} parameter(0) - %constant.865 = bf16[]{:T(256)} constant(-inf) - %pad.40 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.903, %constant.865), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} - %pad.39 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.903, %constant.865), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} - ROOT %maximum.35 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} maximum(%pad.40, %pad.39), metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} -} - -%region_27.32 (reduce_sum.163: f32[], reduce_sum.164: f32[]) -> f32[] { - %reduce_sum.163 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.164 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.168 = f32[]{:T(128)} add(%reduce_sum.163, %reduce_sum.164), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_26.31 (reduce_sum.157: f32[], reduce_sum.161: f32[]) -> f32[] { - %reduce_sum.157 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.161 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.162 = f32[]{:T(128)} add(%reduce_sum.157, %reduce_sum.161), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.329 (param_0.1123: f32[4,4096], param_1.1283: f32[4,4096]) -> (f32[], f32[]) { - %param_0.1123 = f32[4,4096]{1,0:T(4,128)S(1)} parameter(0) - %bitcast.371 = f32[4096,4]{0,1:T(4,128)} bitcast(%param_0.1123), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %square.195 = f32[4096,4]{0,1:T(4,128)} multiply(%bitcast.371, %bitcast.371), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1010 = f32[]{:T(128)} constant(0) - %reduce.150 = f32[]{:T(128)} reduce(%square.195, %constant.1010), dimensions={0,1}, to_apply=%region_27.32, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %param_1.1283 = f32[4,4096]{1,0:T(4,128)S(1)} parameter(1) - %bitcast.375.clone.1 = f32[4096,4]{0,1:T(4,128)} bitcast(%param_1.1283), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %square.198.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%bitcast.375.clone.1, %bitcast.375.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %reduce.151.clone.1 = f32[]{:T(128)} reduce(%square.198.clone.1, %constant.1010), dimensions={0,1}, to_apply=%region_26.31, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.157 = (f32[]{:T(128)}, f32[]{:T(128)}) tuple(%reduce.150, %reduce.151.clone.1) -} - -%region_54.59 (reduce_sum.301: f32[], reduce_sum.302: f32[]) -> f32[] { - %reduce_sum.301 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.302 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.303 = f32[]{:T(128)} add(%reduce_sum.301, %reduce_sum.302), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_40.45 (reduce_sum.226: f32[], reduce_sum.227: f32[]) -> f32[] { - %reduce_sum.226 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.227 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.231 = f32[]{:T(128)} add(%reduce_sum.226, %reduce_sum.227), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.332 (param_0.1112: f32[4096,4], param_1.1275: f32[], param_2.1106: f32[], param_3.795: f32[], param_4.496: f32[4096,4], param_5.421: f32[], param_6.293: f32[4,4096], param_7.192: pred[], param_8.110: f32[4096,4]) -> (f32[], f32[4096,4], f32[4096,4], f32[4096,4], f32[]) { - %param_0.1112 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(0) - %param_3.795 = f32[]{:T(128)S(6)} parameter(3) - %mul.1536.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_3.795), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.192 = pred[]{:T(512)S(6)} parameter(7) - %select_n.266.clone.1 = pred[4096,4]{0,1:T(4,128)(4,1)} broadcast(%param_7.192), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.293 = f32[4,4096]{1,0:T(4,128)S(1)} parameter(6) - %bitcast.419.clone.1 = f32[4096,4]{0,1:T(4,128)} bitcast(%param_6.293), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.421 = f32[]{:T(128)} parameter(5) - %div.773.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_5.421), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.772.clone.1 = f32[4096,4]{0,1:T(4,128)} divide(%bitcast.419.clone.1, %div.773.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.265.clone.1 = f32[4096,4]{0,1:T(4,128)} select(%select_n.266.clone.1, %bitcast.419.clone.1, %div.772.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %constant.943.clone.1 = f32[]{:T(128)} constant(0.1) - %broadcast.578.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%constant.943.clone.1), dimensions={}, metadata={op_name="broadcast.68"} - %mul.1540.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%select_n.265.clone.1, %broadcast.578.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.110 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(8) - %constant.947.clone.1 = f32[]{:T(128)} constant(0.9) - %broadcast.577.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%constant.947.clone.1), dimensions={}, metadata={op_name="broadcast.67"} - %mul.1539.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%param_8.110, %broadcast.577.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.809.clone.1 = f32[4096,4]{0,1:T(4,128)S(1)} add(%mul.1540.clone.1, %mul.1539.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1106 = f32[]{:T(128)S(6)} parameter(2) - %div.769.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_2.1106), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %integer_pow.66.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%select_n.265.clone.1, %select_n.265.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} - %constant.946.clone.1 = f32[]{:T(128)} constant(0.05) - %broadcast.576.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%constant.946.clone.1), dimensions={}, metadata={op_name="broadcast.58"} - %mul.1538.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%integer_pow.66.clone.1, %broadcast.576.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.496 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(4) - %constant.945.clone.1 = f32[]{:T(128)} constant(0.95) - %broadcast.575.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%constant.945.clone.1), dimensions={}, metadata={op_name="broadcast.57"} - %mul.1537.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%param_4.496, %broadcast.575.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.808.clone.1 = f32[4096,4]{0,1:T(4,128)S(1)} add(%mul.1538.clone.1, %mul.1537.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1275 = f32[]{:T(128)S(6)} parameter(1) - %div.768.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_1.1275), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.767.clone.1 = f32[4096,4]{0,1:T(4,128)} divide(%add.808.clone.1, %div.768.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %sqrt.64.clone.1 = f32[4096,4]{0,1:T(4,128)} sqrt(%div.767.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} - %constant.944.clone.1 = f32[]{:T(128)} constant(1e-08) - %broadcast.573.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%constant.944.clone.1), dimensions={}, metadata={op_name="broadcast.53"} - %add.807.clone.1 = f32[4096,4]{0,1:T(4,128)} add(%sqrt.64.clone.1, %broadcast.573.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %multiply.262.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%div.769.clone.1, %add.807.clone.1), metadata={op_name="multiply.36"} - %div.766.clone.1 = f32[4096,4]{0,1:T(4,128)} divide(%add.809.clone.1, %multiply.262.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1535.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%param_0.1112, %broadcast.578.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.806.clone.1 = f32[4096,4]{0,1:T(4,128)} add(%div.766.clone.1, %mul.1535.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1534.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%mul.1536.clone.1, %add.806.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.805.clone.1 = f32[4096,4]{0,1:T(4,128)S(1)} add(%param_0.1112, %mul.1534.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.199 = f32[4096,4]{0,1:T(4,128)} multiply(%add.805.clone.1, %add.805.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.999 = f32[]{:T(128)} constant(0) - %reduce.152 = f32[]{:T(128)} reduce(%square.199, %constant.999), dimensions={0,1}, to_apply=%region_54.59, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.154.clone.1 = f32[]{:T(128)} reduce(%integer_pow.66.clone.1, %constant.999), dimensions={0,1}, to_apply=%region_40.45, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.144 = (f32[]{:T(128)}, f32[4096,4]{0,1:T(4,128)S(1)}, f32[4096,4]{0,1:T(4,128)S(1)}, f32[4096,4]{0,1:T(4,128)S(1)}, f32[]{:T(128)}) tuple(%reduce.152, %add.805.clone.1, %add.808.clone.1, %add.809.clone.1, %reduce.154.clone.1) -} - -%region_53.58 (reduce_sum.295: f32[], reduce_sum.296: f32[]) -> f32[] { - %reduce_sum.295 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.296 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.297 = f32[]{:T(128)} add(%reduce_sum.295, %reduce_sum.296), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_39.44 (reduce_sum.220: f32[], reduce_sum.224: f32[]) -> f32[] { - %reduce_sum.220 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.224 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.225 = f32[]{:T(128)} add(%reduce_sum.220, %reduce_sum.224), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.333 (param_0.1113: f32[4096,4], param_1.1276: f32[], param_2.1107: f32[], param_3.796: f32[], param_4.497: f32[4096,4], param_5.422: f32[], param_6.294: f32[4,4096], param_7.193: pred[], param_8.111: f32[4096,4]) -> (f32[], f32[4096,4], f32[4096,4], f32[4096,4], f32[]) { - %param_0.1113 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(0) - %param_3.796 = f32[]{:T(128)S(6)} parameter(3) - %mul.1543.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_3.796), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.193 = pred[]{:T(512)S(6)} parameter(7) - %select_n.270.clone.1 = pred[4096,4]{0,1:T(4,128)(4,1)} broadcast(%param_7.193), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.294 = f32[4,4096]{1,0:T(4,128)S(1)} parameter(6) - %bitcast.421.clone.1 = f32[4096,4]{0,1:T(4,128)} bitcast(%param_6.294), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.422 = f32[]{:T(128)} parameter(5) - %div.781.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_5.422), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.780.clone.1 = f32[4096,4]{0,1:T(4,128)} divide(%bitcast.421.clone.1, %div.781.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.269.clone.1 = f32[4096,4]{0,1:T(4,128)} select(%select_n.270.clone.1, %bitcast.421.clone.1, %div.780.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %constant.949.clone.1 = f32[]{:T(128)} constant(0.1) - %broadcast.584.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%constant.949.clone.1), dimensions={}, metadata={op_name="broadcast.68"} - %mul.1547.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%select_n.269.clone.1, %broadcast.584.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.111 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(8) - %constant.953.clone.1 = f32[]{:T(128)} constant(0.9) - %broadcast.583.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%constant.953.clone.1), dimensions={}, metadata={op_name="broadcast.67"} - %mul.1546.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%param_8.111, %broadcast.583.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.814.clone.1 = f32[4096,4]{0,1:T(4,128)S(1)} add(%mul.1547.clone.1, %mul.1546.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1107 = f32[]{:T(128)S(6)} parameter(2) - %div.777.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_2.1107), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %integer_pow.67.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%select_n.269.clone.1, %select_n.269.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} - %constant.952.clone.1 = f32[]{:T(128)} constant(0.05) - %broadcast.582.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%constant.952.clone.1), dimensions={}, metadata={op_name="broadcast.58"} - %mul.1545.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%integer_pow.67.clone.1, %broadcast.582.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.497 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(4) - %constant.951.clone.1 = f32[]{:T(128)} constant(0.95) - %broadcast.581.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%constant.951.clone.1), dimensions={}, metadata={op_name="broadcast.57"} - %mul.1544.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%param_4.497, %broadcast.581.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.813.clone.1 = f32[4096,4]{0,1:T(4,128)S(1)} add(%mul.1545.clone.1, %mul.1544.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1276 = f32[]{:T(128)S(6)} parameter(1) - %div.776.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%param_1.1276), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.775.clone.1 = f32[4096,4]{0,1:T(4,128)} divide(%add.813.clone.1, %div.776.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %sqrt.65.clone.1 = f32[4096,4]{0,1:T(4,128)} sqrt(%div.775.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} - %constant.950.clone.1 = f32[]{:T(128)} constant(1e-08) - %broadcast.579.clone.1 = f32[4096,4]{0,1:T(4,128)} broadcast(%constant.950.clone.1), dimensions={}, metadata={op_name="broadcast.53"} - %add.812.clone.1 = f32[4096,4]{0,1:T(4,128)} add(%sqrt.65.clone.1, %broadcast.579.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %multiply.263.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%div.777.clone.1, %add.812.clone.1), metadata={op_name="multiply.35"} - %div.774.clone.1 = f32[4096,4]{0,1:T(4,128)} divide(%add.814.clone.1, %multiply.263.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1542.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%param_0.1113, %broadcast.584.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.811.clone.1 = f32[4096,4]{0,1:T(4,128)} add(%div.774.clone.1, %mul.1542.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1541.clone.1 = f32[4096,4]{0,1:T(4,128)} multiply(%mul.1543.clone.1, %add.811.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.810.clone.1 = f32[4096,4]{0,1:T(4,128)S(1)} add(%param_0.1113, %mul.1541.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.200 = f32[4096,4]{0,1:T(4,128)} multiply(%add.810.clone.1, %add.810.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1000 = f32[]{:T(128)} constant(0) - %reduce.153 = f32[]{:T(128)} reduce(%square.200, %constant.1000), dimensions={0,1}, to_apply=%region_53.58, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.155.clone.1 = f32[]{:T(128)} reduce(%integer_pow.67.clone.1, %constant.1000), dimensions={0,1}, to_apply=%region_39.44, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.145 = (f32[]{:T(128)}, f32[4096,4]{0,1:T(4,128)S(1)}, f32[4096,4]{0,1:T(4,128)S(1)}, f32[4096,4]{0,1:T(4,128)S(1)}, f32[]{:T(128)}) tuple(%reduce.153, %add.810.clone.1, %add.813.clone.1, %add.814.clone.1, %reduce.155.clone.1) +%region_6.9 (reduce_max.6: bf16[], reduce_max.8: bf16[]) -> bf16[] { + %reduce_max.6 = bf16[]{:T(256)} parameter(0), metadata={op_name="jit(train_step)/jvp()/reduce_max"} + %reduce_max.8 = bf16[]{:T(256)} parameter(1), metadata={op_name="jit(train_step)/jvp()/reduce_max"} + ROOT %reduce_max.9 = bf16[]{:T(256)} maximum(%reduce_max.6, %reduce_max.8), metadata={op_name="jit(train_step)/jvp()/reduce_max" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%region_9.12 (reduce_sum.99: f32[], reduce_sum.100: f32[]) -> f32[] { - %reduce_sum.100 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.99 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.101 = f32[]{:T(128)} add(%reduce_sum.99, %reduce_sum.100), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +%fused_computation.423 (param_0.1266: bf16[4096,128256], param_1.1408: bf16[4096], param_2.1092: f32[128], param_3.734: bf16[1,128,4096]) -> (bf16[128], bf16[128,128256]) { + %param_1.1408 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(1) + %param_2.1092 = f32[128]{0:T(128)S(1)} parameter(2) + %param_3.734 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} parameter(3) + %fusion.300.clone.1 = bf16[128,4096]{1,0:T(8,128)(2,1)} fusion(%param_1.1408, %param_2.1092, %param_3.734), kind=kLoop, calls=%fused_computation.358.clone.clone + %param_0.1266 = bf16[4096,128256]{1,0:T(8,128)(2,1)} parameter(0) + %fusion.333 = bf16[4096,128256]{1,0:T(8,128)(2,1)} fusion(%param_0.1266), kind=kLoop, calls=%bitcast_fusion.5 + %convolution.119.clone.1 = bf16[128,128256]{1,0:T(8,128)(2,1)S(1)} convolution(%fusion.300.clone.1, %fusion.333), dim_labels=bf_io->bf, metadata={op_name="jit(train_step)/jvp()/dot_general" stack_frame_id=0} + %constant.400.clone.3 = bf16[]{:T(256)} constant(-inf) + %reduce.175 = bf16[128]{0:T(256)(128)(2,1)S(1)} reduce(%convolution.119.clone.1, %constant.400.clone.3), dimensions={1}, to_apply=%region_6.9, metadata={op_name="jit(train_step)/jvp()/reduce_max" stack_frame_id=0} + ROOT %tuple.152 = (bf16[128]{0:T(256)(128)(2,1)S(1)}, bf16[128,128256]{1,0:T(8,128)(2,1)S(1)}) tuple(%reduce.175, %convolution.119.clone.1) } -%fused_computation.344 (param_0.1127: bf16[4096]) -> f32[] { - %param_0.1127 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(0) - %convert_element_type.1006 = f32[4096]{0:T(1024)} convert(%param_0.1127), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %square.203 = f32[4096]{0:T(1024)} multiply(%convert_element_type.1006, %convert_element_type.1006), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1014 = f32[]{:T(128)} constant(0) - ROOT %reduce.156 = f32[]{:T(128)} reduce(%square.203, %constant.1014), dimensions={0}, to_apply=%region_9.12, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} -} - -%region_49.54 (reduce_sum.274: f32[], reduce_sum.275: f32[]) -> f32[] { - %reduce_sum.274 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.275 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.276 = f32[]{:T(128)} add(%reduce_sum.274, %reduce_sum.275), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_35.40 (reduce_sum.199: f32[], reduce_sum.203: f32[]) -> f32[] { - %reduce_sum.199 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.203 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.204 = f32[]{:T(128)} add(%reduce_sum.199, %reduce_sum.203), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.345 (param_0.1117: f32[4096], param_1.1280: f32[], param_2.1111: f32[], param_3.800: f32[], param_4.501: f32[4096], param_5.426: f32[], param_6.298: bf16[4096], param_7.197: pred[], param_8.115: f32[4096]) -> (f32[], f32[4096], f32[4096], f32[4096], f32[]) { - %param_0.1117 = f32[4096]{0:T(1024)S(1)} parameter(0) - %param_3.800 = f32[]{:T(128)S(6)} parameter(3) - %mul.1574.clone.1 = f32[4096]{0:T(1024)} broadcast(%param_3.800), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.197 = pred[]{:T(512)S(6)} parameter(7) - %select_n.286.clone.1 = pred[4096]{0:T(1024)(128)(4,1)} broadcast(%param_7.197), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.298 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(6) - %convert_element_type.1021.clone.1 = f32[4096]{0:T(1024)} convert(%param_6.298), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %param_5.426 = f32[]{:T(128)} parameter(5) - %div.813.clone.1 = f32[4096]{0:T(1024)} broadcast(%param_5.426), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.812.clone.1 = f32[4096]{0:T(1024)} divide(%convert_element_type.1021.clone.1, %div.813.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.285.clone.1 = f32[4096]{0:T(1024)} select(%select_n.286.clone.1, %convert_element_type.1021.clone.1, %div.812.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %constant.973.clone.1 = f32[]{:T(128)} constant(0.1) - %broadcast.600.clone.1 = f32[4096]{0:T(1024)} broadcast(%constant.973.clone.1), dimensions={}, metadata={op_name="broadcast.72"} - %mul.1580.clone.1 = f32[4096]{0:T(1024)} multiply(%select_n.285.clone.1, %broadcast.600.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.115 = f32[4096]{0:T(1024)S(1)} parameter(8) - %constant.977.clone.1 = f32[]{:T(128)} constant(0.9) - %mul.1581.clone.1 = f32[4096]{0:T(1024)} broadcast(%constant.977.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1579.clone.1 = f32[4096]{0:T(1024)} multiply(%param_8.115, %mul.1581.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.836.clone.1 = f32[4096]{0:T(1024)S(1)} add(%mul.1580.clone.1, %mul.1579.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1111 = f32[]{:T(128)S(6)} parameter(2) - %div.809.clone.1 = f32[4096]{0:T(1024)} broadcast(%param_2.1111), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %integer_pow.71.clone.1 = f32[4096]{0:T(1024)} multiply(%select_n.285.clone.1, %select_n.285.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} - %constant.976.clone.1 = f32[]{:T(128)} constant(0.05) - %mul.1578.clone.1 = f32[4096]{0:T(1024)} broadcast(%constant.976.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1576.clone.1 = f32[4096]{0:T(1024)} multiply(%integer_pow.71.clone.1, %mul.1578.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.501 = f32[4096]{0:T(1024)S(1)} parameter(4) - %constant.975.clone.1 = f32[]{:T(128)} constant(0.95) - %mul.1577.clone.1 = f32[4096]{0:T(1024)} broadcast(%constant.975.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1575.clone.1 = f32[4096]{0:T(1024)} multiply(%param_4.501, %mul.1577.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.835.clone.1 = f32[4096]{0:T(1024)S(1)} add(%mul.1576.clone.1, %mul.1575.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1280 = f32[]{:T(128)S(6)} parameter(1) - %div.808.clone.1 = f32[4096]{0:T(1024)} broadcast(%param_1.1280), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.807.clone.1 = f32[4096]{0:T(1024)} divide(%add.835.clone.1, %div.808.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %sqrt.69.clone.1 = f32[4096]{0:T(1024)} sqrt(%div.807.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} - %constant.974.clone.1 = f32[]{:T(128)} constant(1e-08) - %add.834.clone.1 = f32[4096]{0:T(1024)} broadcast(%constant.974.clone.1), dimensions={}, metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %add.833.clone.1 = f32[4096]{0:T(1024)} add(%sqrt.69.clone.1, %add.834.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %multiply.267.clone.1 = f32[4096]{0:T(1024)} multiply(%div.809.clone.1, %add.833.clone.1), metadata={op_name="multiply.31"} - %div.806.clone.1 = f32[4096]{0:T(1024)} divide(%add.836.clone.1, %multiply.267.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1573.clone.1 = f32[4096]{0:T(1024)} multiply(%param_0.1117, %broadcast.600.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.832.clone.1 = f32[4096]{0:T(1024)} add(%div.806.clone.1, %mul.1573.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1572.clone.1 = f32[4096]{0:T(1024)} multiply(%mul.1574.clone.1, %add.832.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.831.clone.1 = f32[4096]{0:T(1024)S(1)} add(%param_0.1117, %mul.1572.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.204 = f32[4096]{0:T(1024)} multiply(%add.831.clone.1, %add.831.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1004 = f32[]{:T(128)} constant(0) - %reduce.157 = f32[]{:T(128)} reduce(%square.204, %constant.1004), dimensions={0}, to_apply=%region_49.54, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.158.clone.1 = f32[]{:T(128)} reduce(%integer_pow.71.clone.1, %constant.1004), dimensions={0}, to_apply=%region_35.40, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.148 = (f32[]{:T(128)}, f32[4096]{0:T(1024)S(1)}, f32[4096]{0:T(1024)S(1)}, f32[4096]{0:T(1024)S(1)}, f32[]{:T(128)}) tuple(%reduce.157, %add.831.clone.1, %add.835.clone.1, %add.836.clone.1, %reduce.158.clone.1) -} - -%fused_computation.351 (param_0.964: s32[512]) -> s32[1024] { - %constant.801 = s32[] constant(0), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} - %broadcast.539 = s32[1024]{0:T(1024)} broadcast(%constant.801), dimensions={}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} - %param_0.964 = s32[512]{0:T(512)S(1)} parameter(0) - %constant.802 = s32[] constant(2147483647), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} - %pad.41 = s32[1024]{0:T(1024)} pad(%param_0.964, %constant.802), padding=0_512, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} - %constant.800 = s32[] constant(128255), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} - %broadcast.538 = s32[1024]{0:T(1024)} broadcast(%constant.800), dimensions={}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} - ROOT %clamp.1 = s32[1024]{0:T(1024)} clamp(%broadcast.539, %pad.41, %broadcast.538), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} -} - -%fused_computation.352 (param_0.963: s32[4,128]) -> s32[512] { - %param_0.963 = s32[4,128]{1,0:T(4,128)} parameter(0) - %constant.888 = s32[]{:T(128)} constant(0) - %broadcast.546 = s32[4,128]{1,0:T(4,128)} broadcast(%constant.888), dimensions={}, metadata={op_name="broadcast.81"} - %lt.32 = pred[4,128]{1,0:T(4,128)(4,1)} compare(%param_0.963, %broadcast.546), direction=LT, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/lt" stack_frame_id=0} - %constant.875 = s32[]{:T(128)} constant(128256) - %add.760 = s32[4,128]{1,0:T(4,128)} broadcast(%constant.875), dimensions={}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/add" stack_frame_id=0} - %add.748 = s32[4,128]{1,0:T(4,128)} add(%param_0.963, %add.760), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/add" stack_frame_id=0} - %select_n.178 = s32[4,128]{1,0:T(4,128)} select(%lt.32, %add.748, %param_0.963), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/select_n" stack_frame_id=0} - ROOT %bitcast.376 = s32[512]{0:T(512)S(1)} bitcast(%select_n.178), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/broadcast_in_dim" stack_frame_id=0} -} - -%region_61.66 (reduce_sum.345: f32[], reduce_sum.346: f32[]) -> f32[] { - %reduce_sum.345 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} - %reduce_sum.346 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} - ROOT %reduce_sum.330 = f32[]{:T(128)} add(%reduce_sum.345, %reduce_sum.346), metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +%region_7.10 (reduce_sum.93: f32[], reduce_sum.94: f32[]) -> f32[] { + %reduce_sum.93 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} + %reduce_sum.94 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} + ROOT %reduce_sum.95 = f32[]{:T(128)} add(%reduce_sum.93, %reduce_sum.94), metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%region_48.53 (reduce_sum.268: f32[], reduce_sum.269: f32[]) -> f32[] { - %reduce_sum.268 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} - %reduce_sum.269 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} - ROOT %reduce_sum.273 = f32[]{:T(128)} add(%reduce_sum.268, %reduce_sum.269), metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +%fused_computation.316 (param_0.1292: bf16[128,128256], param_1.1427: bf16[128]) -> f32[128] { + %param_0.1292 = bf16[128,128256]{1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.868 = f32[128,128256]{1,0:T(8,128)} convert(%param_0.1292), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} + %bitcast.496 = f32[1,128,128256]{2,1,0:T(8,128)} bitcast(%convert_element_type.868), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} + %param_1.1427 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(1) + %sub.100 = f32[1,128,128256]{2,1,0:T(8,128)} broadcast(%param_1.1427), dimensions={1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %sub.71 = f32[1,128,128256]{2,1,0:T(8,128)} subtract(%bitcast.496, %sub.100), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %exp.56 = f32[1,128,128256]{2,1,0:T(8,128)} exponential(%sub.71), metadata={op_name="jit(train_step)/jvp()/exp" stack_frame_id=0} + %constant.389.clone.27 = f32[]{:T(128)} constant(0) + ROOT %reduce.150 = f32[128]{0:T(128)S(1)} reduce(%exp.56, %constant.389.clone.27), dimensions={0,2}, to_apply=%region_7.10, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} +} + +%fused_computation.422 (param_0.1256: f32[128], param_1.1426: bf16[128]) -> (f32[128], f32[128]) { + %param_0.1256 = f32[128]{0:T(128)S(1)} parameter(0) + %log.23 = f32[128]{0:T(128)S(1)} log(%param_0.1256), metadata={op_name="jit(train_step)/jvp()/log" stack_frame_id=0} + %param_1.1426 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(1) + %reduce_max.18.clone.1 = f32[128]{0:T(128)} convert(%param_1.1426), metadata={op_name="jit(train_step)/jvp()/reduce_max" stack_frame_id=0} + %add.796.clone.1 = f32[128]{0:T(128)} add(%log.23, %reduce_max.18.clone.1), metadata={op_name="jit(train_step)/jvp()/add" stack_frame_id=0} + %constant.389.clone.26 = f32[]{:T(128)} constant(0) + %broadcast.628.clone.1 = f32[128]{0:T(128)} broadcast(%constant.389.clone.26), dimensions={}, metadata={op_name="broadcast.94"} + %mul.1508.clone.1 = f32[128]{0:T(128)} multiply(%add.796.clone.1, %broadcast.628.clone.1), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %constant.377.clone.1.clone.1 = f32[]{:T(128)} constant(1) + %broadcast.625.clone.1 = f32[128]{0:T(128)} broadcast(%constant.377.clone.1.clone.1), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/add" stack_frame_id=0} + %add.785.clone.1 = f32[128]{0:T(128)S(1)} add(%mul.1508.clone.1, %broadcast.625.clone.1), metadata={op_name="jit(train_step)/transpose(jvp())/add" stack_frame_id=0} + ROOT %tuple.150 = (f32[128]{0:T(128)S(1)}, f32[128]{0:T(128)S(1)}) tuple(%log.23, %add.785.clone.1) +} + +%region_44.49 (reduce_sum.262: f32[], reduce_sum.266: f32[]) -> f32[] { + %reduce_sum.262 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} + %reduce_sum.266 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} + ROOT %reduce_sum.267 = f32[]{:T(128)} add(%reduce_sum.262, %reduce_sum.266), metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.353 (param_0.1128: bf16[4,128], param_1.1287: f32[4,128], param_2.1114: f32[4,128], param_3.802: s32[4,128]) -> (f32[], f32[], pred[4,128], f32[4,128]) { - %param_3.802 = s32[4,128]{1,0:T(4,128)S(1)} parameter(3) - %constant.979.clone.1 = s32[]{:T(128)} constant(0) - %broadcast.601.clone.1 = s32[4,128]{1,0:T(4,128)} broadcast(%constant.979.clone.1), dimensions={}, metadata={op_name="broadcast.81"} - %ne.6.clone.1 = pred[4,128]{1,0:T(4,128)(4,1)S(1)} compare(%param_3.802, %broadcast.601.clone.1), direction=NE, metadata={op_name="jit(train_step)/jvp()/ne" stack_frame_id=0} - %param_1.1287 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %log.16 = f32[4,128]{1,0:T(4,128)} log(%param_1.1287), metadata={op_name="jit(train_step)/jvp()/log" stack_frame_id=0} - %param_0.1128 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(0) - %reduce_max.15 = f32[4,128]{1,0:T(4,128)} convert(%param_0.1128), metadata={op_name="jit(train_step)/jvp()/reduce_max" stack_frame_id=0} - %add.762 = f32[4,128]{1,0:T(4,128)} add(%log.16, %reduce_max.15), metadata={op_name="jit(train_step)/jvp()/add" stack_frame_id=0} - %square.207 = f32[4,128]{1,0:T(4,128)} multiply(%add.762, %add.762), metadata={op_name="jit(train_step)/jvp()/square" stack_frame_id=0} - %constant.1016 = f32[]{:T(128)} constant(0) - %broadcast.543 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1016), dimensions={}, metadata={op_name="broadcast.32"} - %mul.1473 = f32[4,128]{1,0:T(4,128)} multiply(%square.207, %broadcast.543), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} - %mul.1465 = f32[4,128]{1,0:T(4,128)} select(%ne.6.clone.1, %mul.1473, %broadcast.543), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} - %reduce.159 = f32[]{:T(128)} reduce(%mul.1465, %constant.1016), dimensions={0,1}, to_apply=%region_61.66, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} - %param_2.1114 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %neg.115.clone.1 = f32[4,128]{1,0:T(4,128)} negate(%param_2.1114), metadata={op_name="jit(train_step)/jvp()/neg" stack_frame_id=0} - %add.749.clone.1 = f32[4,128]{1,0:T(4,128)} add(%neg.115.clone.1, %mul.1473), metadata={op_name="jit(train_step)/jvp()/add" stack_frame_id=0} - %mul.1466.clone.1 = f32[4,128]{1,0:T(4,128)} select(%ne.6.clone.1, %add.749.clone.1, %broadcast.543), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} - %reduce.160.clone.1 = f32[]{:T(128)} reduce(%mul.1466.clone.1, %constant.1016), dimensions={0,1}, to_apply=%region_48.53, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} - %mul.1471.clone.1 = f32[4,128]{1,0:T(4,128)} multiply(%add.762, %broadcast.543), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - %constant.891.clone.1 = f32[]{:T(128)} constant(1) - %add.757.clone.1 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.891.clone.1), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/add" stack_frame_id=0} - %add.750.clone.1 = f32[4,128]{1,0:T(4,128)S(1)} add(%mul.1471.clone.1, %add.757.clone.1), metadata={op_name="jit(train_step)/transpose(jvp())/add" stack_frame_id=0} - ROOT %tuple.149 = (f32[]{:T(128)}, f32[]{:T(128)}, pred[4,128]{1,0:T(4,128)(4,1)S(1)}, f32[4,128]{1,0:T(4,128)S(1)}) tuple(%reduce.159, %reduce.160.clone.1, %ne.6.clone.1, %add.750.clone.1) -} - -%fused_computation.356 (param_0.987: f32[4,128], param_1.1101: f32[4,128]) -> f32[4,128] { - %param_0.987 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %param_1.1101 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %constant.869 = f32[]{:T(128)} constant(0.000244140625) - %broadcast.549 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.869), dimensions={}, metadata={op_name="broadcast.264"} - %div.656 = f32[4,128]{1,0:T(4,128)} multiply(%param_1.1101, %broadcast.549), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/div" stack_frame_id=0} - %constant.867 = f32[]{:T(128)} constant(1e-05) - %add.770 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.867), dimensions={}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/add" stack_frame_id=0} - %add.769 = f32[4,128]{1,0:T(4,128)} add(%div.656, %add.770), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/add" stack_frame_id=0} - %rsqrt.90 = f32[4,128]{1,0:T(4,128)} rsqrt(%add.769), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/rsqrt" stack_frame_id=0} - %div.649 = f32[4,128]{1,0:T(4,128)} divide(%rsqrt.90, %add.769), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/div" stack_frame_id=0} - %constant.864 = f32[]{:T(128)} constant(-0.5) - %mul.1477 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.864), dimensions={}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %mul.1470 = f32[4,128]{1,0:T(4,128)} multiply(%div.649, %mul.1477), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %mul.1469 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.987, %mul.1470), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %constant.863 = f32[]{:T(128)} constant(0.00048828125) - %mul.1476 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.863), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - ROOT %mul.1468 = f32[4,128]{1,0:T(4,128)S(1)} multiply(%mul.1469, %mul.1476), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} +%fused_computation.322 (param_0.1291: bf16[128,128256], param_1.1425: f32[128], param_2.1107: bf16[128], param_3.748: s32[128]) -> f32[128] { + %param_3.748 = s32[128]{0:T(128)S(1)} parameter(3) + %eq.26 = s32[1,128,128256]{2,1,0:T(8,128)} broadcast(%param_3.748), dimensions={1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} + %iota.51 = s32[1,128,128256]{2,1,0:T(8,128)} iota(), iota_dimension=2, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} + %eq.21 = pred[1,128,128256]{2,1,0:T(8,128)(4,1)} compare(%eq.26, %iota.51), direction=EQ, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} + %param_0.1291 = bf16[128,128256]{1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.881 = f32[128,128256]{1,0:T(8,128)} convert(%param_0.1291), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} + %bitcast.508 = f32[1,128,128256]{2,1,0:T(8,128)} bitcast(%convert_element_type.881), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} + %param_2.1107 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(2) + %sub.99 = f32[1,128,128256]{2,1,0:T(8,128)} broadcast(%param_2.1107), dimensions={1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %sub.92 = f32[1,128,128256]{2,1,0:T(8,128)} subtract(%bitcast.508, %sub.99), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %param_1.1425 = f32[128]{0:T(128)S(1)} parameter(1) + %sub.97 = f32[1,128,128256]{2,1,0:T(8,128)} broadcast(%param_1.1425), dimensions={1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %sub.91 = f32[1,128,128256]{2,1,0:T(8,128)} subtract(%sub.92, %sub.97), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %constant.389.clone.25 = f32[]{:T(128)} constant(0) + %broadcast.596 = f32[1,128,128256]{2,1,0:T(8,128)} broadcast(%constant.389.clone.25), dimensions={}, metadata={op_name="broadcast.104"} + %mul.1408 = f32[1,128,128256]{2,1,0:T(8,128)} select(%eq.21, %sub.91, %broadcast.596), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + ROOT %reduce.151 = f32[128]{0:T(128)S(1)} reduce(%mul.1408, %constant.389.clone.25), dimensions={0,2}, to_apply=%region_44.49, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} } %region_0.1 (reduce_sum.67: s32[], reduce_sum.71: s32[]) -> s32[] { %reduce_sum.67 = s32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} %reduce_sum.71 = s32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} - ROOT %reduce_sum.72 = s32[]{:T(128)} add(%reduce_sum.67, %reduce_sum.71), metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[{"indices":["0","2"]}]}} + ROOT %reduce_sum.72 = s32[]{:T(128)} add(%reduce_sum.67, %reduce_sum.71), metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.360 (param_0.1004: pred[4,128]) -> s32[] { - %param_0.1004 = pred[4,128]{1,0:T(4,128)(4,1)S(1)} parameter(0) - %convert_element_type.1013 = s32[4,128]{1,0:T(4,128)} convert(%param_0.1004), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} - %constant.889 = s32[]{:T(128)} constant(0) - ROOT %reduce.161 = s32[]{:T(128)} reduce(%convert_element_type.1013, %constant.889), dimensions={0,1}, to_apply=%region_0.1, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} +%fused_computation.402 (param_0.1166: s32[1,128]) -> s32[] { + %param_0.1166 = s32[1,128]{1,0:T(1,128)S(1)} parameter(0) + %constant.378.clone.4 = s32[]{:T(128)} constant(0) + %broadcast.643 = s32[1,128]{1,0:T(1,128)} broadcast(%constant.378.clone.4), dimensions={}, metadata={op_name="broadcast.85"} + %ne.10 = pred[1,128]{1,0:T(4,128)(4,1)} compare(%param_0.1166, %broadcast.643), direction=NE, metadata={op_name="jit(train_step)/jvp()/ne" stack_frame_id=0} + %convert_element_type.941 = s32[1,128]{1,0:T(1,128)} convert(%ne.10), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} + ROOT %reduce.172 = s32[]{:T(128)} reduce(%convert_element_type.941, %constant.378.clone.4), dimensions={0,1}, to_apply=%region_0.1, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} } -%fused_computation.361 (param_0.989: f32[4,128]) -> f32[4,128] { - %param_0.989 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %constant.870 = f32[]{:T(128)} constant(0.000244140625) - %broadcast.541 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.870), dimensions={}, metadata={op_name="broadcast.264"} - %div.654 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.989, %broadcast.541), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/div" stack_frame_id=0} - %constant.868 = f32[]{:T(128)} constant(1e-05) - %add.759 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.868), dimensions={}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/add" stack_frame_id=0} - %add.756 = f32[4,128]{1,0:T(4,128)} add(%div.654, %add.759), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/add" stack_frame_id=0} - ROOT %rsqrt.88 = f32[4,128]{1,0:T(4,128)S(1)} rsqrt(%add.756), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/rsqrt" stack_frame_id=0} +%region_0.1.clone (reduce_sum.416: s32[], reduce_sum.420: s32[]) -> s32[] { + %reduce_sum.416 = s32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} + %reduce_sum.420 = s32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} + ROOT %reduce_sum.421 = s32[]{:T(128)} add(%reduce_sum.416, %reduce_sum.420), metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.362 (param_0.990: pred[4,128], param_1.1286: f32[]) -> f32[4,128] { - %param_0.990 = pred[4,128]{1,0:T(4,128)(4,1)S(1)} parameter(0) - %param_1.1286 = f32[]{:T(128)S(6)} parameter(1) - %broadcast_in_dim.272 = f32[4,128]{1,0:T(4,128)} broadcast(%param_1.1286), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/broadcast_in_dim" stack_frame_id=0} - %constant.1015 = f32[]{:T(128)} constant(0) - %broadcast.545 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1015), dimensions={}, metadata={op_name="broadcast.32"} - ROOT %mul.1478 = f32[4,128]{1,0:T(4,128)S(1)} select(%param_0.990, %broadcast_in_dim.272, %broadcast.545), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} +%region_58.63 (reduce_sum.345: f32[], reduce_sum.346: f32[]) -> f32[] { + %reduce_sum.345 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} + %reduce_sum.346 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} + ROOT %reduce_sum.330 = f32[]{:T(128)} add(%reduce_sum.345, %reduce_sum.346), metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.364 () -> f32[64] { - %constant.873 = f32[]{:T(128)} constant(500000) - %broadcast.552 = f32[64]{0:T(128)} broadcast(%constant.873), dimensions={}, metadata={op_name="broadcast.255"} - %iota.46 = s32[64]{0:T(128)} iota(), iota_dimension=0, metadata={op_name="jit(train_step)/layers/iota" stack_frame_id=0} - %constant.872 = s32[]{:T(128)} constant(2) - %broadcast.551 = s32[64]{0:T(128)} broadcast(%constant.872), dimensions={}, metadata={op_name="broadcast.256"} - %mul.1479 = s32[64]{0:T(128)} multiply(%iota.46, %broadcast.551), metadata={op_name="jit(train_step)/layers/mul" stack_frame_id=0} - %convert_element_type.1014 = f32[64]{0:T(128)} convert(%mul.1479), metadata={op_name="jit(train_step)/layers/convert_element_type" stack_frame_id=0} - %constant.871 = f32[]{:T(128)} constant(0.0078125) - %broadcast.550 = f32[64]{0:T(128)} broadcast(%constant.871), dimensions={}, metadata={op_name="broadcast.257"} - %div.657 = f32[64]{0:T(128)} multiply(%convert_element_type.1014, %broadcast.550), metadata={op_name="jit(train_step)/layers/div" stack_frame_id=0} - ROOT %pow.36 = f32[64]{0:T(128)S(1)} power(%broadcast.552, %div.657), metadata={op_name="jit(train_step)/layers/pow" stack_frame_id=0} +%region_45.50 (reduce_sum.268: f32[], reduce_sum.269: f32[]) -> f32[] { + %reduce_sum.268 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} + %reduce_sum.269 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} + ROOT %reduce_sum.273 = f32[]{:T(128)} add(%reduce_sum.268, %reduce_sum.269), metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.365 (param_0.1002: s32[4,128]) -> (f32[4,128,1,1], f32[4,128]) { - %param_0.1002 = s32[4,128]{1,0:T(4,128)} parameter(0) - %convert_element_type.1015 = f32[4,128]{1,0:T(4,128)S(1)} convert(%param_0.1002), metadata={op_name="jit(train_step)/layers/convert_element_type" stack_frame_id=0} - %bitcast.377 = f32[4,128,1,1]{1,0,3,2:T(4,128)} bitcast(%convert_element_type.1015), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %tuple.151 = (f32[4,128,1,1]{1,0,3,2:T(4,128)}, f32[4,128]{1,0:T(4,128)S(1)}) tuple(%bitcast.377, %convert_element_type.1015) +%fused_computation.405 (param_0.1290: s32[1,128], param_1.1424: bf16[128], param_2.1106: f32[128], param_3.747: f32[128], param_4.464: f32[]) -> (f32[], f32[], f32[128]) { + %param_0.1290 = s32[1,128]{1,0:T(1,128)S(1)} parameter(0) + %constant.378.clone.1 = s32[]{:T(128)} constant(0) + %broadcast.641 = s32[1,128]{1,0:T(1,128)} broadcast(%constant.378.clone.1), dimensions={}, metadata={op_name="broadcast.85"} + %ne.16 = pred[1,128]{1,0:T(4,128)(4,1)} compare(%param_0.1290, %broadcast.641), direction=NE, metadata={op_name="jit(train_step)/jvp()/ne" stack_frame_id=0} + %param_2.1106 = f32[128]{0:T(128)S(1)} parameter(2) + %log.18 = f32[128]{0:T(128)} log(%param_2.1106), metadata={op_name="jit(train_step)/jvp()/log" stack_frame_id=0} + %param_1.1424 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(1) + %reduce_max.16 = f32[128]{0:T(128)} convert(%param_1.1424), metadata={op_name="jit(train_step)/jvp()/reduce_max" stack_frame_id=0} + %add.794 = f32[128]{0:T(128)} add(%log.18, %reduce_max.16), metadata={op_name="jit(train_step)/jvp()/add" stack_frame_id=0} + %square.208 = f32[128]{0:T(128)} multiply(%add.794, %add.794), metadata={op_name="jit(train_step)/jvp()/square" stack_frame_id=0} + %constant.389.clone.24 = f32[]{:T(128)} constant(0) + %broadcast.626 = f32[128]{0:T(128)} broadcast(%constant.389.clone.24), dimensions={}, metadata={op_name="broadcast.94"} + %mul.1516 = f32[128]{0:T(128)} multiply(%square.208, %broadcast.626), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + %bitcast.561 = f32[1,128]{1,0:T(1,128)} bitcast(%mul.1516), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + %broadcast.637 = f32[1,128]{1,0:T(1,128)} broadcast(%constant.389.clone.24), dimensions={}, metadata={op_name="broadcast.94"} + %mul.1502 = f32[1,128]{1,0:T(1,128)} select(%ne.16, %bitcast.561, %broadcast.637), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + %reduce.174 = f32[]{:T(128)} reduce(%mul.1502, %constant.389.clone.24), dimensions={0,1}, to_apply=%region_58.63, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} + %param_3.747 = f32[128]{0:T(128)S(1)} parameter(3) + %neg.118.clone.1 = f32[128]{0:T(128)} negate(%param_3.747), metadata={op_name="jit(train_step)/jvp()/neg" stack_frame_id=0} + %bitcast.565.clone.1 = f32[1,128]{1,0:T(1,128)} bitcast(%neg.118.clone.1), metadata={op_name="jit(train_step)/jvp()/neg" stack_frame_id=0} + %add.784.clone.1 = f32[1,128]{1,0:T(1,128)} add(%bitcast.565.clone.1, %bitcast.561), metadata={op_name="jit(train_step)/jvp()/add" stack_frame_id=0} + %mul.1500.clone.1 = f32[1,128]{1,0:T(1,128)} select(%ne.16, %add.784.clone.1, %broadcast.637), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + %reduce.173.clone.1 = f32[]{:T(128)} reduce(%mul.1500.clone.1, %constant.389.clone.24), dimensions={0,1}, to_apply=%region_45.50, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} + %param_4.464 = f32[]{:T(128)S(6)} parameter(4) + %broadcast_in_dim.304.clone.1 = f32[1,128]{1,0:T(1,128)} broadcast(%param_4.464), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/broadcast_in_dim" stack_frame_id=0} + %mul.1498.clone.1 = f32[1,128]{1,0:T(1,128)} select(%ne.16, %broadcast_in_dim.304.clone.1, %broadcast.637), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %bitcast.559.clone.1 = f32[128]{0:T(128)S(1)} bitcast(%mul.1498.clone.1), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + ROOT %tuple.149 = (f32[]{:T(128)}, f32[]{:T(128)}, f32[128]{0:T(128)S(1)}) tuple(%reduce.174, %reduce.173.clone.1, %bitcast.559.clone.1) +} + +%fused_computation.326.clone.1.clone.clone (param_0.1263: bf16[128,128256], param_1.1407: f32[128], param_2.1084: f32[128], param_3.732: bf16[128], param_4.450: s32[128], param_5.396: f32[128]) -> bf16[128,128256] { + %param_5.396 = f32[128]{0:T(128)S(1)} parameter(5) + %mul.1652 = f32[1,128,128256]{2,1,0:T(8,128)} broadcast(%param_5.396), dimensions={1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %param_1.1407 = f32[128]{0:T(128)S(1)} parameter(1) + %mul.1651 = f32[1,128,128256]{2,1,0:T(8,128)} broadcast(%param_1.1407), dimensions={1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %param_0.1263 = bf16[128,128256]{1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.966 = f32[128,128256]{1,0:T(8,128)} convert(%param_0.1263), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} + %bitcast.653 = f32[1,128,128256]{2,1,0:T(8,128)} bitcast(%convert_element_type.966), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} + %param_3.732 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(3) + %sub.114 = f32[1,128,128256]{2,1,0:T(8,128)} broadcast(%param_3.732), dimensions={1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %sub.113 = f32[1,128,128256]{2,1,0:T(8,128)} subtract(%bitcast.653, %sub.114), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %exp.72 = f32[1,128,128256]{2,1,0:T(8,128)} exponential(%sub.113), metadata={op_name="jit(train_step)/jvp()/exp" stack_frame_id=0} + %mul.1650 = f32[1,128,128256]{2,1,0:T(8,128)} multiply(%mul.1651, %exp.72), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %param_2.1084 = f32[128]{0:T(128)S(1)} parameter(2) + %div.833 = f32[1,128,128256]{2,1,0:T(8,128)} broadcast(%param_2.1084), dimensions={1}, metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} + %div.832 = f32[1,128,128256]{2,1,0:T(8,128)} divide(%mul.1650, %div.833), metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} + %param_4.450 = s32[128]{0:T(128)S(1)} parameter(4) + %eq.35 = s32[1,128,128256]{2,1,0:T(8,128)} broadcast(%param_4.450), dimensions={1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} + %iota.60 = s32[1,128,128256]{2,1,0:T(8,128)} iota(), iota_dimension=2, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} + %eq.34 = pred[1,128,128256]{2,1,0:T(8,128)(4,1)} compare(%eq.35, %iota.60), direction=EQ, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} + %convert_element_type.965 = f32[1,128,128256]{2,1,0:T(8,128)} convert(%eq.34), metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/convert_element_type" stack_frame_id=0} + %sub.112 = f32[1,128,128256]{2,1,0:T(8,128)} subtract(%div.832, %convert_element_type.965), metadata={op_name="jit(train_step)/transpose(jvp())/sub" stack_frame_id=0} + %mul.1649 = f32[1,128,128256]{2,1,0:T(8,128)} multiply(%mul.1652, %sub.112), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %convert_element_type.964 = bf16[1,128,128256]{2,1,0:T(8,128)(2,1)} convert(%mul.1649), metadata={op_name="jit(train_step)/transpose(jvp())/convert_element_type" stack_frame_id=0} + ROOT %bitcast.652 = bf16[128,128256]{1,0:T(8,128)(2,1)} bitcast(%convert_element_type.964) +} + +%bitcast_fusion.4 (bitcast_input.4: bf16[4096,128256]) -> bf16[4096,128256] { + %bitcast_input.4 = bf16[4096,128256]{1,0:T(8,128)(2,1)} parameter(0) + ROOT %bitcast.660 = bf16[4096,128256]{1,0:T(8,128)(2,1)} bitcast(%bitcast_input.4) } -%fused_computation.369 (param_0.1103: f32[4096,4]) -> bf16[4,4096] { - %param_0.1103 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(0) - %bitcast.451 = f32[4,4096]{1,0:T(4,128)} bitcast(%param_0.1103), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - ROOT %convert.106 = bf16[4,4096]{1,0:T(4,128)(2,1)} convert(%bitcast.451) +%region_8.11 (dot_general.182: bf16[], dot_general.183: bf16[]) -> bf16[] { + %dot_general.182 = bf16[]{:T(256)} parameter(0), metadata={op_name="jit(train_step)/transpose(jvp(...k,k->...k))/dot_general"} + %dot_general.183 = bf16[]{:T(256)} parameter(1), metadata={op_name="jit(train_step)/transpose(jvp(...k,k->...k))/dot_general"} + ROOT %add.165 = bf16[]{:T(256)} add(%dot_general.182, %dot_general.183), metadata={op_name="add.51"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%fused_computation.359 (param_0.1262: f32[128], param_1.1406: bf16[1,128,4096], param_2.1085: bf16[4096,128256], param_3.733: bf16[128,128256], param_4.451: f32[128], param_5.397: f32[128], param_6.252: bf16[128], param_7.168: s32[128], param_8.97: f32[128]) -> (bf16[4096], bf16[1,128,4096]) { + %param_3.733 = bf16[128,128256]{1,0:T(8,128)(2,1)S(1)} parameter(3) + %param_4.451 = f32[128]{0:T(128)S(1)} parameter(4) + %param_5.397 = f32[128]{0:T(128)S(1)} parameter(5) + %param_6.252 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(6) + %param_7.168 = s32[128]{0:T(128)S(1)} parameter(7) + %param_8.97 = f32[128]{0:T(128)S(1)} parameter(8) + %fusion.320.clone.1 = bf16[128,128256]{1,0:T(8,128)(2,1)} fusion(%param_3.733, %param_4.451, %param_5.397, %param_6.252, %param_7.168, /*index=5*/%param_8.97), kind=kLoop, calls=%fused_computation.326.clone.1.clone.clone + %param_2.1085 = bf16[4096,128256]{1,0:T(8,128)(2,1)} parameter(2) + %fusion.332 = bf16[4096,128256]{1,0:T(8,128)(2,1)} fusion(%param_2.1085), kind=kLoop, calls=%bitcast_fusion.4 + %convolution.117.clone.1 = bf16[128,4096]{1,0:T(8,128)(2,1)} convolution(%fusion.320.clone.1, %fusion.332), dim_labels=bf_oi->bf, metadata={op_name="jit(train_step)/transpose(jvp())/dot_general" stack_frame_id=0} + %bitcast.545.clone.1 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%convolution.117.clone.1), metadata={op_name="jit(train_step)/transpose(jvp())/dot_general" stack_frame_id=0} + %convert.130 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%bitcast.545.clone.1) + %param_1.1406 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %convert_element_type.920 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%param_1.1406), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} + %param_0.1262 = f32[128]{0:T(128)S(1)} parameter(0) + %mul.1463 = f32[1,128,4096]{2,1,0:T(8,128)} broadcast(%param_0.1262), dimensions={1}, metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + %mul.1462 = f32[1,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.920, %mul.1463), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + %convert_element_type.919 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1462), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} + %convert.131 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%convert_element_type.919) + %multiply.271 = f32[1,128,4096]{2,1,0:T(8,128)} multiply(%convert.130, %convert.131), metadata={op_name="multiply.197"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","4096"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %convert.132 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} convert(%multiply.271) + %constant.392.clone.1 = bf16[]{:T(256)} constant(0) + %reduce.162 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} reduce(%convert.132, %constant.392.clone.1), dimensions={0,1}, to_apply=%region_8.11, metadata={op_name="jit(train_step)/transpose(jvp(...k,k->...k))/dot_general" stack_frame_id=0} + ROOT %tuple.153 = (bf16[4096]{0:T(1024)(128)(2,1)S(1)}, bf16[1,128,4096]{2,1,0:T(8,128)(2,1)S(1)}) tuple(%reduce.162, %bitcast.545.clone.1) +} + +%region_8.11.clone (dot_general.243: bf16[], dot_general.244: bf16[]) -> bf16[] { + %dot_general.243 = bf16[]{:T(256)} parameter(0), metadata={op_name="jit(train_step)/transpose(jvp(...k,k->...k))/dot_general"} + %dot_general.244 = bf16[]{:T(256)} parameter(1), metadata={op_name="jit(train_step)/transpose(jvp(...k,k->...k))/dot_general"} + ROOT %add.500 = bf16[]{:T(256)} add(%dot_general.243, %dot_general.244), metadata={op_name="add.51"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.370 (param_0.1104: f32[4096,4]) -> bf16[4,4096] { - %param_0.1104 = f32[4096,4]{0,1:T(4,128)S(1)} parameter(0) - %bitcast.452 = f32[4,4096]{1,0:T(4,128)} bitcast(%param_0.1104), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - ROOT %convert.108 = bf16[4,4096]{1,0:T(4,128)(2,1)S(1)} convert(%bitcast.452) +%region_9.12 (reduce_sum.99: f32[], reduce_sum.100: f32[]) -> f32[] { + %reduce_sum.99 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} + %reduce_sum.100 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} + ROOT %reduce_sum.101 = f32[]{:T(128)} add(%reduce_sum.99, %reduce_sum.100), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%region_6.9 (reduce_max.6: bf16[], reduce_max.8: bf16[]) -> bf16[] { - %reduce_max.6 = bf16[]{:T(256)} parameter(0), metadata={op_name="jit(train_step)/jvp()/reduce_max"} - %reduce_max.8 = bf16[]{:T(256)} parameter(1), metadata={op_name="jit(train_step)/jvp()/reduce_max"} - ROOT %reduce_max.9 = bf16[]{:T(256)} maximum(%reduce_max.6, %reduce_max.8), metadata={op_name="jit(train_step)/jvp()/reduce_max" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +%fused_computation.388 (param_0.1289: bf16[4096]) -> f32[] { + %param_0.1289 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(0) + %convert_element_type.932 = f32[4096]{0:T(1024)} convert(%param_0.1289), metadata={op_name="jit(train_step)/transpose(jvp())/convert_element_type" stack_frame_id=0} + %square.204 = f32[4096]{0:T(1024)} multiply(%convert_element_type.932, %convert_element_type.932), metadata={op_name="jit(train_step)/square" stack_frame_id=0} + %constant.389.clone.23 = f32[]{:T(128)} constant(0) + ROOT %reduce.169 = f32[]{:T(128)} reduce(%square.204, %constant.389.clone.23), dimensions={0}, to_apply=%region_9.12, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} } -%fused_computation.237.clone.clone (param_0.1090: f32[4096,128256]) -> bf16[4096,128256,1] { - %param_0.1090 = f32[4096,128256]{1,0:T(8,128)} parameter(0) - %convert_element_type.1026 = bf16[4096,128256]{1,0:T(8,128)(2,1)} convert(%param_0.1090), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} - ROOT %bitcast.447 = bf16[4096,128256,1]{1,0,2:T(8,128)(2,1)} bitcast(%convert_element_type.1026), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} +%region_10.13 (reduce_sum.102: f32[], reduce_sum.106: f32[]) -> f32[] { + %reduce_sum.102 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/transpose(jvp())/reduce_sum"} + %reduce_sum.106 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/transpose(jvp())/reduce_sum"} + ROOT %reduce_sum.107 = f32[]{:T(128)} add(%reduce_sum.102, %reduce_sum.106), metadata={op_name="jit(train_step)/transpose(jvp())/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%fused_computation.354 (param_0.1288: bf16[1,128,4096], param_1.1423: bf16[1,128,4096], param_2.1105: bf16[4096]) -> f32[128] { + %param_0.1288 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.908 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%param_0.1288), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} + %param_1.1423 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %convert.128 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%param_1.1423) + %param_2.1105 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(2) + %dot_general.397 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1105), dimensions={2}, metadata={op_name="jit(train_step)/jvp(...k,k->...k)/dot_general" stack_frame_id=0} + %convert.129 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%dot_general.397) + %dot_general.383 = f32[1,128,4096]{2,1,0:T(8,128)} multiply(%convert.128, %convert.129), metadata={op_name="jit(train_step)/transpose(jvp(...k,k->...k))/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","4096"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %convert_element_type.907 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%dot_general.383), metadata={op_name="jit(train_step)/transpose(jvp())/convert_element_type" stack_frame_id=0} + %mul.1448 = f32[1,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.908, %convert_element_type.907), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %constant.389.clone.22 = f32[]{:T(128)} constant(0) + ROOT %reduce.161 = f32[128]{0:T(128)S(1)} reduce(%mul.1448, %constant.389.clone.22), dimensions={0,2}, to_apply=%region_10.13, metadata={op_name="jit(train_step)/transpose(jvp())/reduce_sum" stack_frame_id=0} +} + +%fused_computation.414 (param_0.1160: f32[128], param_1.1244: f32[128]) -> f32[128] { + %param_0.1160 = f32[128]{0:T(128)S(1)} parameter(0) + %bitcast.575 = f32[1,128]{1,0:T(1,128)} bitcast(%param_0.1160), metadata={op_name="jit(train_step)/transpose(jvp())/reduce_sum" stack_frame_id=0} + %param_1.1244 = f32[128]{0:T(128)S(1)} parameter(1) + %constant.398.clone.1 = f32[]{:T(128)} constant(0.000244140625) + %broadcast.631 = f32[128]{0:T(128)} broadcast(%constant.398.clone.1), dimensions={}, metadata={op_name="broadcast.283"} + %div.665 = f32[128]{0:T(128)} multiply(%param_1.1244, %broadcast.631), metadata={op_name="jit(train_step)/jvp()/div" stack_frame_id=0} + %constant.399.clone.1 = f32[]{:T(128)} constant(1e-05) + %broadcast.629 = f32[128]{0:T(128)} broadcast(%constant.399.clone.1), dimensions={}, metadata={op_name="jit(train_step)/jvp()/add" stack_frame_id=0} + %add.789 = f32[128]{0:T(128)} add(%div.665, %broadcast.629), metadata={op_name="jit(train_step)/jvp()/add" stack_frame_id=0} + %bitcast.574 = f32[1,128]{1,0:T(1,128)} bitcast(%add.789), metadata={op_name="jit(train_step)/jvp()/add" stack_frame_id=0} + %rsqrt.107 = f32[1,128]{1,0:T(1,128)} rsqrt(%bitcast.574), metadata={op_name="jit(train_step)/jvp()/rsqrt" stack_frame_id=0} + %div.663 = f32[1,128]{1,0:T(1,128)} divide(%rsqrt.107, %bitcast.574), metadata={op_name="jit(train_step)/jvp()/div" stack_frame_id=0} + %constant.401.clone.1 = f32[]{:T(128)} constant(-0.5) + %mul.1520 = f32[1,128]{1,0:T(1,128)} broadcast(%constant.401.clone.1), dimensions={}, metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + %mul.1514 = f32[1,128]{1,0:T(1,128)} multiply(%div.663, %mul.1520), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + %mul.1513 = f32[1,128]{1,0:T(1,128)} multiply(%bitcast.575, %mul.1514), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %constant.402.clone.1 = f32[]{:T(128)} constant(0.00048828125) + %mul.1519 = f32[1,128]{1,0:T(1,128)} broadcast(%constant.402.clone.1), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %mul.1509 = f32[1,128]{1,0:T(1,128)} multiply(%mul.1513, %mul.1519), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + ROOT %bitcast.569 = f32[128]{0:T(128)S(1)} bitcast(%mul.1509), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} +} + +%fused_computation.351 (param_0.1002: bf16[1,128,4096], param_1.1053: f32[128], param_2.700: f32[128], param_3.393: bf16[1,128,4096], param_4.231: bf16[4096]) -> bf16[1,128,4096] { + %param_3.393 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) + %convert.126 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%param_3.393) + %param_4.231 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(4) + %dot_general.398 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_4.231), dimensions={2}, metadata={op_name="jit(train_step)/jvp(...k,k->...k)/dot_general" stack_frame_id=0} + %convert.127 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%dot_general.398) + %dot_general.384 = f32[1,128,4096]{2,1,0:T(8,128)} multiply(%convert.126, %convert.127), metadata={op_name="jit(train_step)/transpose(jvp(...k,k->...k))/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","4096"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %convert_element_type.899 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%dot_general.384), metadata={op_name="jit(train_step)/transpose(jvp())/convert_element_type" stack_frame_id=0} + %param_2.700 = f32[128]{0:T(128)S(1)} parameter(2) + %mul.1452 = f32[1,128,4096]{2,1,0:T(8,128)} broadcast(%param_2.700), dimensions={1}, metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + %mul.1444 = f32[1,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.899, %mul.1452), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %param_0.1002 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.910 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%param_0.1002), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} + %param_1.1053 = f32[128]{0:T(128)S(1)} parameter(1) + %mul.1451 = f32[1,128,4096]{2,1,0:T(8,128)} broadcast(%param_1.1053), dimensions={1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %mul.1450 = f32[1,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.910, %mul.1451), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %add_any.153 = f32[1,128,4096]{2,1,0:T(8,128)} add(%mul.1444, %mul.1450), metadata={op_name="jit(train_step)/transpose(jvp())/add_any" stack_frame_id=0} + ROOT %convert_element_type.897 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} convert(%add_any.153), metadata={op_name="jit(train_step)/transpose(jvp())/convert_element_type" stack_frame_id=0} +} + +%fused_computation.84.clone.1 (param_0.1453: bf16[4,1024,14336], param_1.1527: s32[]) -> bf16[1,1024,14336] { + %param_0.1453 = bf16[4,1024,14336]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1527 = s32[]{:T(128)S(6)} parameter(1) + %constant.362.clone.72 = s32[]{:T(128)} constant(0) + ROOT %dynamic-slice.175 = bf16[1,1024,14336]{2,1,0:T(8,128)(2,1)S(1)} dynamic-slice(%param_0.1453, %param_1.1527, %constant.362.clone.72, %constant.362.clone.72), dynamic_slice_sizes={1,1024,14336}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/jit(dynamic_index_in_dim)/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} +} + +%convert_element_type.523.reduce_sub_computation (lhs: bf16[], rhs: bf16[]) -> bf16[] { + %lhs = bf16[] parameter(0) + %rhs = bf16[] parameter(1) + ROOT %add.659 = bf16[] add(%lhs, %rhs), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.317.clone.clone (param_0.1091: f32[4,128], param_1.1257: bf16[4,128,4096], param_2.1077: bf16[4096]) -> bf16[4,128,4096] { - %param_2.1077 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(2) - %dot_general.383 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1077), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1257 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %convert_element_type.1028 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_1.1257), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %param_0.1091 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.1595 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_0.1091), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %mul.1594 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1028, %mul.1595), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %convert_element_type.1027 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1594), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - ROOT %dot_general.382 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.383, %convert_element_type.1027), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} +%fused_computation.239.clone.1 (param_0.1427: bf16[4,4096], param_1.1511: s32[]) -> bf16[4096] { + %param_0.1427 = bf16[4,4096]{1,0:T(4,128)(2,1)} parameter(0) + %param_1.1511 = s32[]{:T(128)S(6)} parameter(1) + %constant.362.clone.63 = s32[]{:T(128)} constant(0) + %dynamic_slice.289 = bf16[1,4096]{1,0:T(2,128)(2,1)} dynamic-slice(%param_0.1427, %param_1.1511, %constant.362.clone.63), dynamic_slice_sizes={1,4096}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/jit(dynamic_index_in_dim)/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %constant.996 = bf16[]{:T(256)} constant(-0), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + ROOT %reduce.188 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} reduce(%dynamic_slice.289, %constant.996), dimensions={0}, to_apply=%convert_element_type.523.reduce_sub_computation, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} } -%fused_computation.371 (param_0.1105: f32[4096,128256], param_1.1268: f32[4,128], param_2.1099: bf16[4,128,4096], param_3.788: bf16[4096]) -> (bf16[4,128], bf16[4,128,128256]) { - %param_1.1268 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %param_2.1099 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) - %param_3.788 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(3) - %fusion.240.clone.1 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} fusion(%param_1.1268, %param_2.1099, %param_3.788), kind=kLoop, calls=%fused_computation.317.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %param_0.1105 = f32[4096,128256]{1,0:T(8,128)} parameter(0) - %fusion.221.clone.1 = bf16[4096,128256,1]{1,0,2:T(8,128)(2,1)} fusion(%param_0.1105), kind=kLoop, calls=%fused_computation.237.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/convert_element_type" stack_frame_id=0} - %convolution.87.clone.1 = bf16[4,128,128256]{2,1,0:T(8,128)(2,1)} convolution(%fusion.240.clone.1, %fusion.221.clone.1), window={size=1}, dim_labels=0bf_io0->0bf, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/logits_dense/dot_general" stack_frame_id=0} - %constant.992 = bf16[]{:T(256)} constant(-inf) - %reduce.162 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} reduce(%convolution.87.clone.1, %constant.992), dimensions={2}, to_apply=%region_6.9, metadata={op_name="jit(train_step)/jvp()/reduce_max" stack_frame_id=0} - ROOT %tuple.152 = (bf16[4,128]{1,0:T(4,128)(2,1)S(1)}, bf16[4,128,128256]{2,1,0:T(8,128)(2,1)}) tuple(%reduce.162, %convolution.87.clone.1) +%fused_computation.98.clone.1 (param_0.1448: bf16[4,32,128,1024], param_1.1525: s32[]) -> bf16[1,32,128,1024] { + %param_0.1448 = bf16[4,32,128,1024]{3,2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1525 = s32[]{:T(128)S(6)} parameter(1) + %constant.362.clone.71 = s32[]{:T(128)} constant(0) + ROOT %dynamic-slice.174 = bf16[1,32,128,1024]{3,2,1,0:T(8,128)(2,1)S(1)} dynamic-slice(%param_0.1448, %param_1.1525, %constant.362.clone.71, %constant.362.clone.71, %constant.362.clone.71), dynamic_slice_sizes={1,32,128,1024}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/jit(dynamic_index_in_dim)/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} } -%fused_computation.372 (param_0.1102: f32[4096,4,8,128]) -> bf16[4,4096,8,128] { - %param_0.1102 = f32[4096,4,8,128]{3,2,1,0:T(8,128)} parameter(0) - %bitcast.450 = f32[4,4096,8,128]{3,2,0,1:T(8,128)} bitcast(%param_0.1102), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - ROOT %convert.110 = bf16[4,4096,8,128]{3,2,0,1:T(8,128)(2,1)} convert(%bitcast.450) +%fused_computation.99.clone.1 (param_0.1442: bf16[4,1024,32,128], param_1.1521: s32[]) -> bf16[1,1024,32,128] { + %param_0.1442 = bf16[4,1024,32,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) + %param_1.1521 = s32[]{:T(128)S(6)} parameter(1) + %constant.362.clone.69 = s32[]{:T(128)} constant(0) + ROOT %dynamic-slice.172 = bf16[1,1024,32,128]{1,3,2,0:T(8,128)(2,1)S(1)} dynamic-slice(%param_0.1442, %param_1.1521, %constant.362.clone.69, %constant.362.clone.69, %constant.362.clone.69), dynamic_slice_sizes={1,1024,32,128}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/jit(dynamic_index_in_dim)/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} } %convert_element_type.525.reduce_sub_computation (lhs.1: bf16[], rhs.1: bf16[]) -> bf16[] { %lhs.1 = bf16[] parameter(0) %rhs.1 = bf16[] parameter(1) - ROOT %add.624 = bf16[] add(%lhs.1, %rhs.1), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + ROOT %add.660 = bf16[] add(%lhs.1, %rhs.1), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.121.clone.clone (param_0.1242: bf16[4,4096], param_1.1376: s32[]) -> bf16[4096] { - %param_0.1242 = bf16[4,4096]{1,0:T(4,128)(2,1)} parameter(0) - %param_1.1376 = s32[]{:T(128)S(6)} parameter(1) - %constant.1116 = s32[]{:T(128)} constant(0) - %dynamic_slice.316 = bf16[1,4096]{1,0:T(2,128)(2,1)} dynamic-slice(%param_0.1242, %param_1.1376, %constant.1116), dynamic_slice_sizes={1,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - %constant.1117 = bf16[]{:T(256)} constant(-0), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %reduce.174 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} reduce(%dynamic_slice.316, %constant.1117), dimensions={0}, to_apply=%convert_element_type.525.reduce_sub_computation, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} +%fused_computation.238.clone.1 (param_0.1428: bf16[4,4096], param_1.1512: s32[]) -> bf16[4096] { + %param_0.1428 = bf16[4,4096]{1,0:T(4,128)(2,1)} parameter(0) + %param_1.1512 = s32[]{:T(128)S(6)} parameter(1) + %constant.362.clone.64 = s32[]{:T(128)} constant(0) + %dynamic_slice.290 = bf16[1,4096]{1,0:T(2,128)(2,1)} dynamic-slice(%param_0.1428, %param_1.1512, %constant.362.clone.64), dynamic_slice_sizes={1,4096}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/jit(dynamic_index_in_dim)/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %constant.997 = bf16[]{:T(256)} constant(-0), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + ROOT %reduce.189 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} reduce(%dynamic_slice.290, %constant.997), dimensions={0}, to_apply=%convert_element_type.525.reduce_sub_computation, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} } %region_12.14 (reduce_sum.108: f32[], reduce_sum.109: f32[]) -> f32[] { - %reduce_sum.108 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/reduce_sum"} - %reduce_sum.109 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/reduce_sum"} - ROOT %reduce_sum.113 = f32[]{:T(128)} add(%reduce_sum.108, %reduce_sum.109), metadata={op_name="checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.58.clone.clone (param_0.1243: bf16[4,4,128,4096], param_1.1377: s32[]) -> f32[4,128] { - %param_0.1243 = bf16[4,4,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1377 = s32[]{:T(128)S(6)} parameter(1) - %constant.1118 = s32[]{:T(128)} constant(0) - %dynamic_slice.317 = bf16[1,4,128,4096]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1243, %param_1.1377, %constant.1118, %constant.1118, %constant.1118), dynamic_slice_sizes={1,4,128,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - %bitcast.548 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.317), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} - %convert_element_type.1093 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%bitcast.548), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %square.214 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1093, %convert_element_type.1093), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/square" stack_frame_id=0} - %constant.1119 = f32[]{:T(128)} constant(0) - ROOT %reduce.175 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%square.214, %constant.1119), dimensions={2}, to_apply=%region_12.14, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0} -} - -%fused_computation.143.clone.1.clone (param_0.1244: f32[4,128]) -> f32[4,128] { - %param_0.1244 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %constant.1121 = f32[]{:T(128)} constant(0.000244140625) - %closed_call.81 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1121), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} - %div.842 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.1244, %closed_call.81), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} - %constant.1120 = f32[]{:T(128)} constant(1e-05) - %closed_call.80 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1120), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} - %add.858 = f32[4,128]{1,0:T(4,128)} add(%div.842, %closed_call.80), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} - ROOT %rsqrt.97 = f32[4,128]{1,0:T(4,128)S(1)} rsqrt(%add.858), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/rsqrt" stack_frame_id=0} -} - -%fused_computation.24.clone.1.clone.clone (param_0.1258: bf16[4,4096,32,128], param_1.1387: s32[]) -> bf16[4096,32,128,1] { - %param_0.1258 = bf16[4,4096,32,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) - %param_1.1387 = s32[]{:T(128)S(6)} parameter(1) - %constant.1134 = s32[]{:T(128)} constant(0) - %dynamic_slice.323 = bf16[1,4096,32,128]{1,3,2,0:T(8,128)(2,1)} dynamic-slice(%param_0.1258, %param_1.1387, %constant.1134, %constant.1134, %constant.1134), dynamic_slice_sizes={1,4096,32,128}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.559 = bf16[4096,32,128,1]{0,2,1,3:T(8,128)(2,1)} bitcast(%dynamic_slice.323), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} -} - -%fused_computation.91.clone.clone (param_0.1259: f32[4,128], param_1.1388: bf16[4,4,128,4096], param_2.1176: s32[], param_3.847: bf16[4096]) -> bf16[4,128,4096,1] { - %param_3.847 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(3) - %dot_general.428 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_3.847), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1388 = bf16[4,4,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(1) - %param_2.1176 = s32[]{:T(128)S(6)} parameter(2) - %constant.1135 = s32[]{:T(128)} constant(0) - %dynamic_slice.324 = bf16[1,4,128,4096]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_1.1388, %param_2.1176, %constant.1135, %constant.1135, %constant.1135), dynamic_slice_sizes={1,4,128,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - %bitcast.561 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.324), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} - %convert_element_type.1101 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%bitcast.561), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_0.1259 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.1709 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_0.1259), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.1708 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1101, %mul.1709), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1100 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1708), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %dot_general.427 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.428, %convert_element_type.1100), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - ROOT %bitcast.560 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} bitcast(%dot_general.427), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} -} - -%fused_computation.36.clone.clone (param_0.1260: bf16[4,4096,32,128], param_1.1389: s32[], param_2.1177: f32[4,128], param_3.848: bf16[4,4,128,4096], param_4.530: bf16[4096]) -> bf16[4,128,32,128] { - %param_2.1177 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %param_3.848 = bf16[4,4,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(3) - %param_1.1389 = s32[]{:T(128)S(6)} parameter(1) - %param_4.530 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(4) - %fusion.343 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} fusion(%param_2.1177, %param_3.848, %param_1.1389, %param_4.530), kind=kLoop, calls=%fused_computation.91.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_0.1260 = bf16[4,4096,32,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) - %fusion.342 = bf16[4096,32,128,1]{0,2,1,3:T(8,128)(2,1)} fusion(%param_0.1260, %param_1.1389), kind=kLoop, calls=%fused_computation.24.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %convolution.113 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)S(1)} convolution(%fusion.343, %fusion.342), window={size=1x32 pad=0_0x31_31 rhs_reversal=0x1}, dim_labels=0bf1_i1o0->0b1f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} -} - -%fused_computation.70.clone.clone (param_0.1261: bf16[4,128,32,128]) -> (bf16[4,128,32,64], bf16[4,128,32,64]) { - %param_0.1261 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) - %split.160 = bf16[4,128,32,64]{3,1,2,0:T(8,128)(2,1)} slice(%param_0.1261), slice={[0:4], [0:128], [0:32], [64:128]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/split" stack_frame_id=0} - %neg.129 = bf16[4,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)} negate(%split.160), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/neg" stack_frame_id=0} - %split.161 = bf16[4,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)} slice(%param_0.1261), slice={[0:4], [0:128], [0:32], [0:64]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/split" stack_frame_id=0} - ROOT %tuple.187 = (bf16[4,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)}, bf16[4,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)}) tuple(%neg.129, %split.161) -} - -%fused_computation.145.clone.clone () -> f32[64] { - %constant.1124 = f32[]{:T(128)} constant(500000) - %closed_call.84 = f32[64]{0:T(128)} broadcast(%constant.1124), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} - %iota.51 = s32[64]{0:T(128)} iota(), iota_dimension=0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/iota" stack_frame_id=0} - %constant.1123 = s32[]{:T(128)} constant(2) - %closed_call.83 = s32[64]{0:T(128)} broadcast(%constant.1123), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} - %mul.1699 = s32[64]{0:T(128)} multiply(%iota.51, %closed_call.83), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1094 = f32[64]{0:T(128)} convert(%mul.1699), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %constant.1122 = f32[]{:T(128)} constant(0.0078125) - %closed_call.82 = f32[64]{0:T(128)} broadcast(%constant.1122), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} - %div.843 = f32[64]{0:T(128)} multiply(%convert_element_type.1094, %closed_call.82), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} - ROOT %pow.38 = f32[64]{0:T(128)S(1)} power(%closed_call.84, %div.843), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/pow" stack_frame_id=0} -} - -%fused_computation.117.clone.clone (param_0.1245: f32[64], param_1.1378: f32[4,128]) -> (bf16[4,128,1,64], bf16[4,128,1,64]) { - %param_1.1378 = f32[4,128]{1,0:T(4,128)} parameter(1) - %div.846 = f32[4,128,1,64]{3,1,0,2:T(8,128)} broadcast(%param_1.1378), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} - %param_0.1245 = f32[64]{0:T(128)S(1)} parameter(0) - %div.845 = f32[4,128,1,64]{3,1,0,2:T(8,128)} broadcast(%param_0.1245), dimensions={3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} - %div.844 = f32[4,128,1,64]{3,1,0,2:T(8,128)} divide(%div.846, %div.845), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} - %cos.43 = f32[4,128,1,64]{3,1,0,2:T(8,128)} cosine(%div.844), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/cos" stack_frame_id=0} - %convert_element_type.1095 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} convert(%cos.43), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %sin.35.clone.3 = f32[4,128,1,64]{3,1,0,2:T(8,128)} sine(%div.844), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/sin" stack_frame_id=0} - %convert_element_type.829.clone.3 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} convert(%sin.35.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %tuple.185 = (bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)}, bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)}) tuple(%convert_element_type.1095, %convert_element_type.829.clone.3) -} - -%fused_computation.120.clone.clone (param_0.1252: bf16[4,128,1,64]) -> bf16[4,128,128] { - %param_0.1252 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} parameter(0) - %constant.1130 = bf16[]{:T(256)} constant(-inf) - %pad.61 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1252, %constant.1130), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - %pad.60 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1252, %constant.1130), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - %maximum.45 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} maximum(%pad.61, %pad.60), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - ROOT %bitcast.554 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%maximum.45), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} -} - -%fused_computation.119.clone.clone (param_0.1246: bf16[4,128,1,64]) -> bf16[4,128,128] { - %param_0.1246 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} parameter(0) - %constant.1125 = bf16[]{:T(256)} constant(-inf) - %pad.59 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1246, %constant.1125), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - %pad.58 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1246, %constant.1125), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - %maximum.44 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} maximum(%pad.59, %pad.58), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - ROOT %bitcast.549 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%maximum.44), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} -} - -%fused_computation.73.clone.clone (param_0.1262: bf16[4,128,32,64], param_1.1390: bf16[4,128,32,64], param_2.1178: bf16[4,128,32,128], param_3.849: bf16[4,128,128], param_4.531: bf16[4,128,128]) -> bf16[4,32,128,128] { - %param_2.1178 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(2) - %param_4.531 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(4) - %mul.1713 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_4.531), dimensions={0,1,3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.1711 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} multiply(%param_2.1178, %mul.1713), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %param_1.1390 = bf16[4,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) - %constant.1136 = bf16[]{:T(256)} constant(-inf) - %pad.65 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_1.1390, %constant.1136), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - %param_0.1262 = bf16[4,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) - %pad.64 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1262, %constant.1136), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - %maximum.47 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} maximum(%pad.65, %pad.64), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - %param_3.849 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) - %mul.1712 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_3.849), dimensions={0,1,3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.1710 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} multiply(%maximum.47, %mul.1712), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %add.860 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} add(%mul.1711, %mul.1710), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} - ROOT %bitcast.562 = bf16[4,32,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} bitcast(%add.860), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} -} - -%fused_computation.90.clone.clone (param_0.1254: f32[4,128], param_1.1384: bf16[4,4,128,4096], param_2.1173: s32[], param_3.844: bf16[4096]) -> bf16[4,128,4096,1] { - %param_3.844 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(3) - %dot_general.426 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_3.844), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1384 = bf16[4,4,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(1) - %param_2.1173 = s32[]{:T(128)S(6)} parameter(2) - %constant.1132 = s32[]{:T(128)} constant(0) - %dynamic_slice.322 = bf16[1,4,128,4096]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_1.1384, %param_2.1173, %constant.1132, %constant.1132, %constant.1132), dynamic_slice_sizes={1,4,128,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - %bitcast.557 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.322), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} - %convert_element_type.1099 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%bitcast.557), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_0.1254 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.1703 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_0.1254), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.1702 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1099, %mul.1703), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1098 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1702), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %dot_general.425 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.426, %convert_element_type.1098), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - ROOT %bitcast.556 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} bitcast(%dot_general.425), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} -} - -%fused_computation.64.clone.1.clone.clone (param_0.1253: bf16[4,4096,8,128], param_1.1383: s32[]) -> bf16[4096,8,128,1] { - %param_0.1253 = bf16[4,4096,8,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) - %param_1.1383 = s32[]{:T(128)S(6)} parameter(1) - %constant.1131 = s32[]{:T(128)} constant(0) - %dynamic_slice.321 = bf16[1,4096,8,128]{1,3,2,0:T(8,128)(2,1)} dynamic-slice(%param_0.1253, %param_1.1383, %constant.1131, %constant.1131, %constant.1131), dynamic_slice_sizes={1,4096,8,128}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.555 = bf16[4096,8,128,1]{0,2,1,3:T(8,128)(2,1)} bitcast(%dynamic_slice.321), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} -} - -%fused_computation.89.clone.clone (param_0.1255: bf16[4,4096,8,128], param_1.1385: s32[], param_2.1174: f32[4,128], param_3.845: bf16[4,4,128,4096], param_4.528: bf16[4096]) -> bf16[4,128,8,128] { - %param_2.1174 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %param_3.845 = bf16[4,4,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(3) - %param_1.1385 = s32[]{:T(128)S(6)} parameter(1) - %param_4.528 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(4) - %fusion.340 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} fusion(%param_2.1174, %param_3.845, %param_1.1385, %param_4.528), kind=kLoop, calls=%fused_computation.90.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_0.1255 = bf16[4,4096,8,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) - %fusion.341 = bf16[4096,8,128,1]{0,2,1,3:T(8,128)(2,1)} fusion(%param_0.1255, %param_1.1385), kind=kLoop, calls=%fused_computation.64.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %convolution.112 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} convolution(%fusion.340, %fusion.341), window={size=1x8 pad=0_0x7_7 rhs_reversal=0x1}, dim_labels=0bf1_i1o0->0b1f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} -} - -%fused_computation.106.clone.clone (param_0.1256: bf16[4,128,8,128]) -> (bf16[4,128,8,64], bf16[4,128,8,64]) { - %param_0.1256 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) - %split.158 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)} slice(%param_0.1256), slice={[0:4], [0:128], [0:8], [64:128]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/split" stack_frame_id=0} - %neg.128 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} negate(%split.158), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/neg" stack_frame_id=0} - %split.159 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} slice(%param_0.1256), slice={[0:4], [0:128], [0:8], [0:64]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/split" stack_frame_id=0} - ROOT %tuple.186 = (bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)}, bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)}) tuple(%neg.128, %split.159) -} - -%fused_computation.109.clone.clone (param_0.1257: bf16[4,128,8,64], param_1.1386: bf16[4,128,8,64], param_2.1175: bf16[4,128,8,128], param_3.846: bf16[4,128,128], param_4.529: bf16[4,128,128]) -> bf16[4,8,128,128] { - %param_2.1175 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(2) - %param_4.529 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(4) - %mul.1707 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_4.529), dimensions={0,1,3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.1705 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} multiply(%param_2.1175, %mul.1707), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %param_1.1386 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) - %constant.1133 = bf16[]{:T(256)} constant(-inf) - %pad.63 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_1.1386, %constant.1133), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - %param_0.1257 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) - %pad.62 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1257, %constant.1133), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - %maximum.46 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} maximum(%pad.63, %pad.62), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - %param_3.846 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) - %mul.1706 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_3.846), dimensions={0,1,3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.1704 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} multiply(%maximum.46, %mul.1706), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %add.859 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} add(%mul.1705, %mul.1704), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} - ROOT %bitcast.558 = bf16[4,8,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} bitcast(%add.859), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} -} - -%fused_computation.135.clone.clone (param_0.1248: bf16[4,4096,8,128], param_1.1380: s32[]) -> bf16[1,4096,8,128] { - %param_0.1248 = bf16[4,4096,8,128]{3,2,0,1:T(8,128)(2,1)} parameter(0) - %param_1.1380 = s32[]{:T(128)S(6)} parameter(1) - %constant.1128 = s32[]{:T(128)} constant(0) - ROOT %dynamic_slice.319 = bf16[1,4096,8,128]{3,2,0,1:T(8,128)(2,1)S(1)} dynamic-slice(%param_0.1248, %param_1.1380, %constant.1128, %constant.1128, %constant.1128), dynamic_slice_sizes={1,4096,8,128}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} -} - -%fused_computation.65.clone.1.clone.clone.clone.clone (param_0.1249: bf16[1,4096,8,128]) -> bf16[4096,8,128,1] { - %param_0.1249 = bf16[1,4096,8,128]{3,2,0,1:T(8,128)(2,1)S(1)} parameter(0) - %copy.248 = bf16[1,4096,8,128]{3,1,2,0:T(8,128)(2,1)} copy(%param_0.1249), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0} - ROOT %bitcast.550 = bf16[4096,8,128,1]{2,0,1,3:T(8,128)(2,1)} bitcast(%copy.248), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} -} - -%fused_computation.88.clone.clone.clone.clone (param_0.1250: f32[4,128], param_1.1381: bf16[4,4,128,4096], param_2.1171: s32[], param_3.842: bf16[4096]) -> bf16[4,128,4096,1] { - %param_3.842 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(3) - %dot_general.424 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_3.842), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1381 = bf16[4,4,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(1) - %param_2.1171 = s32[]{:T(128)S(6)} parameter(2) - %constant.1129 = s32[]{:T(128)} constant(0) - %dynamic_slice.320 = bf16[1,4,128,4096]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_1.1381, %param_2.1171, %constant.1129, %constant.1129, %constant.1129), dynamic_slice_sizes={1,4,128,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - %bitcast.552 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.320), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} - %convert_element_type.1097 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%bitcast.552), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_0.1250 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.1701 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_0.1250), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.1700 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1097, %mul.1701), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1096 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1700), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %dot_general.423 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.424, %convert_element_type.1096), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - ROOT %bitcast.551 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} bitcast(%dot_general.423), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} -} - -%fused_computation.114.clone.clone (param_0.1251: bf16[1,4096,8,128], param_1.1382: f32[4,128], param_2.1172: bf16[4,4,128,4096], param_3.843: s32[], param_4.527: bf16[4096]) -> bf16[4,8,128,128] { - %param_1.1382 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %param_2.1172 = bf16[4,4,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(2) - %param_3.843 = s32[]{:T(128)S(6)} parameter(3) - %param_4.527 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(4) - %fusion.339 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} fusion(%param_1.1382, %param_2.1172, %param_3.843, %param_4.527), kind=kLoop, calls=%fused_computation.88.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_0.1251 = bf16[1,4096,8,128]{3,2,0,1:T(8,128)(2,1)S(1)} parameter(0) - %fusion.338 = bf16[4096,8,128,1]{2,0,1,3:T(8,128)(2,1)} fusion(%param_0.1251), kind=kLoop, calls=%fused_computation.65.clone.1.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %convolution.111 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} convolution(%fusion.339, %fusion.338), window={size=1x8 pad=0_0x7_7 rhs_reversal=0x1}, dim_labels=0bf1_i1o0->0b1f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} - ROOT %bitcast.553 = bf16[4,8,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} bitcast(%convolution.111), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} -} - -%fused_computation.366.clone.clone (param_0.1286: f32[4,32,128,128]) -> (f32[4,32,128,1], f32[4,32,128]) { - %param_0.1286 = f32[4,32,128,128]{2,1,0,3:T(8,128)S(1)} parameter(0) - %slice.11 = f32[4,32,128,1]{2,1,0,3:T(8,128)S(1)} slice(%param_0.1286), slice={[0:4], [0:32], [0:128], [0:1]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/shard_map/vmap(jit(_splash_attention))/slice" stack_frame_id=0} - %bitcast.262.clone.3 = f32[4,32,128]{2,1,0:T(8,128)S(1)} bitcast(%slice.11), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/shard_map/vmap(jit(_splash_attention))/squeeze" stack_frame_id=0} - ROOT %tuple.192 = (f32[4,32,128,1]{2,1,0,3:T(8,128)S(1)}, f32[4,32,128]{2,1,0:T(8,128)S(1)}) tuple(%slice.11, %bitcast.262.clone.3) + %reduce_sum.108 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/reduce_sum"} + %reduce_sum.109 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/reduce_sum"} + ROOT %reduce_sum.113 = f32[]{:T(128)} add(%reduce_sum.108, %reduce_sum.109), metadata={op_name="checkpoint/rematted_computation/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%fused_computation.153.clone.1 (param_0.1425: bf16[4,1,128,4096], param_1.1510: s32[]) -> f32[128] { + %param_0.1425 = bf16[4,1,128,4096]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %param_1.1510 = s32[]{:T(128)S(6)} parameter(1) + %constant.362.clone.62 = s32[]{:T(128)} constant(0) + %dynamic-slice.167 = bf16[1,1,128,4096]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1425, %param_1.1510, %constant.362.clone.62, %constant.362.clone.62, %constant.362.clone.62), dynamic_slice_sizes={1,1,128,4096}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/jit(dynamic_index_in_dim)/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %bitcast.784 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic-slice.167), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/jit(dynamic_index_in_dim)/squeeze" stack_frame_id=0} + %convert_element_type.1015 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%bitcast.784), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %square.215 = f32[1,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1015, %convert_element_type.1015), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/square" stack_frame_id=0} + %constant.363.clone.11 = f32[]{:T(128)} constant(0) + ROOT %reduce.187 = f32[128]{0:T(128)S(1)} reduce(%square.215, %constant.363.clone.11), dimensions={0,2}, to_apply=%region_12.14, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/reduce_sum" stack_frame_id=0} +} + +%fused_computation.243.clone.1 (param_0.1426: f32[128]) -> f32[128] { + %param_0.1426 = f32[128]{0:T(128)S(1)} parameter(0) + %constant.364.clone.9 = f32[]{:T(128)} constant(0.000244140625) + %broadcast.709 = f32[128]{0:T(128)} broadcast(%constant.364.clone.9), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call" stack_frame_id=0} + %div.852 = f32[128]{0:T(128)} multiply(%param_0.1426, %broadcast.709), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/div" stack_frame_id=0} + %constant.365.clone.9 = f32[]{:T(128)} constant(1e-05) + %broadcast.708 = f32[128]{0:T(128)} broadcast(%constant.365.clone.9), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call" stack_frame_id=0} + %add.886 = f32[128]{0:T(128)} add(%div.852, %broadcast.708), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/add" stack_frame_id=0} + %bitcast.786 = f32[1,128]{1,0:T(1,128)} bitcast(%add.886), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/add" stack_frame_id=0} + %rsqrt.114 = f32[1,128]{1,0:T(1,128)} rsqrt(%bitcast.786), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/rsqrt" stack_frame_id=0} + ROOT %bitcast.785 = f32[128]{0:T(128)S(1)} bitcast(%rsqrt.114), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/rsqrt" stack_frame_id=0} +} + +%fused_computation.194.clone.clone.clone.1 (param_0.1444: bf16[4096], param_1.1522: f32[128], param_2.1166: bf16[4,1,128,4096], param_3.788: s32[]) -> bf16[128,4096,1] { + %param_0.1444 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(0) + %dot_general.447 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_0.1444), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/...k,k->...k/dot_general" stack_frame_id=0} + %convert.151 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%dot_general.447) + %param_2.1166 = bf16[4,1,128,4096]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %param_3.788 = s32[]{:T(128)S(6)} parameter(3) + %constant.362.clone.70 = s32[]{:T(128)} constant(0) + %dynamic-slice.173 = bf16[1,1,128,4096]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_2.1166, %param_3.788, %constant.362.clone.70, %constant.362.clone.70, %constant.362.clone.70), dynamic_slice_sizes={1,1,128,4096}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/jit(dynamic_index_in_dim)/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %bitcast.800 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic-slice.173), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/jit(dynamic_index_in_dim)/squeeze" stack_frame_id=0} + %convert_element_type.1023 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%bitcast.800), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %param_1.1522 = f32[128]{0:T(128)S(1)} parameter(1) + %mul.1746 = f32[1,128,4096]{2,1,0:T(8,128)} broadcast(%param_1.1522), dimensions={1}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %mul.1745 = f32[1,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1023, %mul.1746), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %convert_element_type.1022 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1745), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %convert.152 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%convert_element_type.1022) + %dot_general.446 = f32[1,128,4096]{2,1,0:T(8,128)} multiply(%convert.151, %convert.152), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/...k,k->...k/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","4096"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %convert.153 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} convert(%dot_general.446) + ROOT %bitcast.799 = bf16[128,4096,1]{1,0,2:T(8,128)(2,1)} bitcast(%convert.153) +} + +%fused_computation.94.clone.clone.clone.1 (param_0.1443: bf16[1,4096,32,128]) -> bf16[4096,32,128] { + %param_0.1443 = bf16[1,4096,32,128]{1,3,2,0:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.798 = bf16[4096,32,128]{0,2,1:T(8,128)(2,1)} bitcast(%param_0.1443), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/dot_general" stack_frame_id=0} +} + +%fused_computation.201.clone.1 (param_0.1445: bf16[1,4096,32,128], param_1.1523: bf16[4096], param_2.1167: f32[128], param_3.789: bf16[4,1,128,4096], param_4.487: s32[]) -> bf16[1,128,32,128] { + %param_1.1523 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(1) + %param_2.1167 = f32[128]{0:T(128)S(1)} parameter(2) + %param_3.789 = bf16[4,1,128,4096]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(3) + %param_4.487 = s32[]{:T(128)S(6)} parameter(4) + %fusion.426 = bf16[128,4096,1]{1,0,2:T(8,128)(2,1)} fusion(%param_1.1523, %param_2.1167, %param_3.789, %param_4.487), kind=kLoop, calls=%fused_computation.194.clone.clone.clone.1 + %param_0.1445 = bf16[1,4096,32,128]{1,3,2,0:T(8,128)(2,1)S(1)} parameter(0) + %fusion.425 = bf16[4096,32,128]{0,2,1:T(8,128)(2,1)} fusion(%param_0.1445), kind=kLoop, calls=%fused_computation.94.clone.clone.clone.1, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/dot_general" stack_frame_id=0} + %convolution.144 = bf16[128,32,128]{2,0,1:T(8,128)(2,1)} convolution(%fusion.426, %fusion.425), window={size=32 pad=31_31 rhs_reversal=1}, dim_labels=bf0_i0o->b0f, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/dot_general" stack_frame_id=0} + ROOT %bitcast.801 = bf16[1,128,32,128]{3,1,2,0:T(8,128)(2,1)S(1)} bitcast(%convolution.144), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/dot_general" stack_frame_id=0} +} + +%fused_computation.165.clone.1 (param_0.1446: bf16[1,128,32,128]) -> (bf16[1,128,32,64], bf16[1,128,32,64]) { + %param_0.1446 = bf16[1,128,32,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %slice.75 = bf16[1,128,32,64]{3,1,2,0:T(8,128)(2,1)} slice(%param_0.1446), slice={[0:1], [0:128], [0:32], [64:128]}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/split" stack_frame_id=0} + %convert.197 = f32[1,128,32,64]{3,1,2,0:T(8,128)} convert(%slice.75) + %neg.132 = f32[1,128,32,64]{3,1,2,0:T(8,128)} negate(%convert.197), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/neg" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","32","64"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %convert.198 = bf16[1,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)} convert(%neg.132) + %slice.76 = bf16[1,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)} slice(%param_0.1446), slice={[0:1], [0:128], [0:32], [0:64]}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/split" stack_frame_id=0} + ROOT %tuple.186 = (bf16[1,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)}, bf16[1,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)}) tuple(%convert.198, %slice.76) +} + +%fused_computation.252.clone.1 () -> f32[64] { + %constant.366.clone.3 = f32[]{:T(128)} constant(500000) + %closed_call.64 = f32[64]{0:T(128)} broadcast(%constant.366.clone.3), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call" stack_frame_id=0} + %iota.65 = s32[64]{0:T(128)} iota(), iota_dimension=0, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/iota" stack_frame_id=0} + %constant.367.clone.3 = s32[]{:T(128)} constant(2) + %closed_call.52 = s32[64]{0:T(128)} broadcast(%constant.367.clone.3), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call" stack_frame_id=0} + %mul.1738 = s32[64]{0:T(128)} multiply(%iota.65, %closed_call.52), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %convert_element_type.1018 = f32[64]{0:T(128)} convert(%mul.1738), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %constant.368.clone.3 = f32[]{:T(128)} constant(0.0078125) + %closed_call.51 = f32[64]{0:T(128)} broadcast(%constant.368.clone.3), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call" stack_frame_id=0} + %div.853 = f32[64]{0:T(128)} multiply(%convert_element_type.1018, %closed_call.51), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/div" stack_frame_id=0} + ROOT %pow.38 = f32[64]{0:T(128)S(1)} power(%closed_call.64, %div.853), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/pow" stack_frame_id=0} +} + +%fused_computation.234.clone.1 (param_0.1433: f32[128], param_1.1516: f32[64]) -> (bf16[1,128,1,64], bf16[1,128,1,64]) { + %param_0.1433 = f32[128]{0:T(128)S(1)} parameter(0) + %div.855 = f32[1,128,1,64]{3,1,2,0:T(8,128)} broadcast(%param_0.1433), dimensions={1}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/div" stack_frame_id=0} + %param_1.1516 = f32[64]{0:T(128)S(1)} parameter(1) + %div.856 = f32[1,128,1,64]{3,1,2,0:T(8,128)} broadcast(%param_1.1516), dimensions={3}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/div" stack_frame_id=0} + %div.854 = f32[1,128,1,64]{3,1,2,0:T(8,128)} divide(%div.855, %div.856), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/div" stack_frame_id=0} + %cos.43 = f32[1,128,1,64]{3,1,2,0:T(8,128)} cosine(%div.854), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/cos" stack_frame_id=0} + %convert_element_type.1019 = bf16[1,128,1,64]{3,1,2,0:T(8,128)(2,1)S(1)} convert(%cos.43), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %sin.35.clone.3 = f32[1,128,1,64]{3,1,2,0:T(8,128)} sine(%div.854), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/sin" stack_frame_id=0} + %convert_element_type.768.clone.3 = bf16[1,128,1,64]{3,1,2,0:T(8,128)(2,1)S(1)} convert(%sin.35.clone.3), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + ROOT %tuple.184 = (bf16[1,128,1,64]{3,1,2,0:T(8,128)(2,1)S(1)}, bf16[1,128,1,64]{3,1,2,0:T(8,128)(2,1)S(1)}) tuple(%convert_element_type.1019, %convert_element_type.768.clone.3) +} + +%fused_computation.237.clone.1 (param_0.1435: bf16[1,128,1,64]) -> bf16[128,128] { + %param_0.1435 = bf16[1,128,1,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %constant.999 = bf16[]{:T(256)} constant(-inf) + %pad.61 = bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1435, %constant.999), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/concatenate" stack_frame_id=0} + %convert.199 = f32[1,128,1,128]{3,1,2,0:T(8,128)} convert(%pad.61) + %pad.60 = bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1435, %constant.999), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/concatenate" stack_frame_id=0} + %convert.200 = f32[1,128,1,128]{3,1,2,0:T(8,128)} convert(%pad.60) + %maximum.49 = f32[1,128,1,128]{3,1,2,0:T(8,128)} maximum(%convert.199, %convert.200), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/concatenate" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","1","128"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %convert.201 = bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)} convert(%maximum.49) + ROOT %bitcast.792 = bf16[128,128]{1,0:T(8,128)(2,1)S(1)} bitcast(%convert.201), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} +} + +%fused_computation.236.clone.1 (param_0.1434: bf16[1,128,1,64]) -> bf16[128,128] { + %param_0.1434 = bf16[1,128,1,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %constant.998 = bf16[]{:T(256)} constant(-inf) + %pad.59 = bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1434, %constant.998), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/concatenate" stack_frame_id=0} + %convert.202 = f32[1,128,1,128]{3,1,2,0:T(8,128)} convert(%pad.59) + %pad.58 = bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1434, %constant.998), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/concatenate" stack_frame_id=0} + %convert.203 = f32[1,128,1,128]{3,1,2,0:T(8,128)} convert(%pad.58) + %maximum.48 = f32[1,128,1,128]{3,1,2,0:T(8,128)} maximum(%convert.202, %convert.203), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/concatenate" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","1","128"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %convert.204 = bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)} convert(%maximum.48) + ROOT %bitcast.791 = bf16[128,128]{1,0:T(8,128)(2,1)S(1)} bitcast(%convert.204), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} +} + +%fused_computation.168.clone.1 (param_0.1447: bf16[1,128,32,64], param_1.1524: bf16[1,128,32,64], param_2.1168: bf16[1,128,32,128], param_3.790: bf16[128,128], param_4.488: bf16[128,128]) -> bf16[32,128,128] { + %param_2.1168 = bf16[1,128,32,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(2) + %convert.205 = f32[1,128,32,128]{3,1,2,0:T(8,128)} convert(%param_2.1168) + %param_4.488 = bf16[128,128]{1,0:T(8,128)(2,1)S(1)} parameter(4) + %mul.1750 = bf16[1,128,32,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_4.488), dimensions={1,3}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %convert.206 = f32[1,128,32,128]{3,1,2,0:T(8,128)} convert(%mul.1750) + %mul.1748 = f32[1,128,32,128]{3,1,2,0:T(8,128)} multiply(%convert.205, %convert.206), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","32","128"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %param_1.1524 = bf16[1,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) + %constant.1001 = bf16[]{:T(256)} constant(-inf) + %pad.65 = bf16[1,128,32,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_1.1524, %constant.1001), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/concatenate" stack_frame_id=0} + %convert.207 = f32[1,128,32,128]{3,1,2,0:T(8,128)} convert(%pad.65) + %param_0.1447 = bf16[1,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %pad.64 = bf16[1,128,32,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1447, %constant.1001), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/concatenate" stack_frame_id=0} + %convert.208 = f32[1,128,32,128]{3,1,2,0:T(8,128)} convert(%pad.64) + %maximum.51 = f32[1,128,32,128]{3,1,2,0:T(8,128)} maximum(%convert.207, %convert.208), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/concatenate" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","32","128"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %param_3.790 = bf16[128,128]{1,0:T(8,128)(2,1)S(1)} parameter(3) + %mul.1749 = bf16[1,128,32,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_3.790), dimensions={1,3}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %convert.209 = f32[1,128,32,128]{3,1,2,0:T(8,128)} convert(%mul.1749) + %mul.1747 = f32[1,128,32,128]{3,1,2,0:T(8,128)} multiply(%maximum.51, %convert.209), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","32","128"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %add.888 = f32[1,128,32,128]{3,1,2,0:T(8,128)} add(%mul.1748, %mul.1747), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/add" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","32","128"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %convert.210 = bf16[1,128,32,128]{3,1,2,0:T(8,128)(2,1)} convert(%add.888) + ROOT %bitcast.802 = bf16[32,128,128]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%convert.210), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/shard_map/vmap(static_sparsity_splash)/jit(_splash_attention)/splash_mha_fwd_segmented_residuals/splash_mha_fwd_segmented_residuals/squeeze" stack_frame_id=0} +} + +%fused_computation.117.clone.1 (param_0.1436: bf16[4,1024,8,128], param_1.1517: s32[]) -> bf16[1,1024,8,128] { + %param_0.1436 = bf16[4,1024,8,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) + %param_1.1517 = s32[]{:T(128)S(6)} parameter(1) + %constant.362.clone.67 = s32[]{:T(128)} constant(0) + ROOT %dynamic-slice.170 = bf16[1,1024,8,128]{1,3,2,0:T(8,128)(2,1)S(1)} dynamic-slice(%param_0.1436, %param_1.1517, %constant.362.clone.67, %constant.362.clone.67, %constant.362.clone.67), dynamic_slice_sizes={1,1024,8,128}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/jit(dynamic_index_in_dim)/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} +} + +%fused_computation.193.clone.clone.clone.1 (param_0.1438: bf16[4096], param_1.1518: f32[128], param_2.1163: bf16[4,1,128,4096], param_3.785: s32[]) -> bf16[128,4096,1] { + %param_0.1438 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(0) + %dot_general.445 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_0.1438), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/...k,k->...k/dot_general" stack_frame_id=0} + %convert.148 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%dot_general.445) + %param_2.1163 = bf16[4,1,128,4096]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %param_3.785 = s32[]{:T(128)S(6)} parameter(3) + %constant.362.clone.68 = s32[]{:T(128)} constant(0) + %dynamic-slice.171 = bf16[1,1,128,4096]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_2.1163, %param_3.785, %constant.362.clone.68, %constant.362.clone.68, %constant.362.clone.68), dynamic_slice_sizes={1,1,128,4096}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/jit(dynamic_index_in_dim)/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %bitcast.795 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic-slice.171), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/jit(dynamic_index_in_dim)/squeeze" stack_frame_id=0} + %convert_element_type.1021 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%bitcast.795), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %param_1.1518 = f32[128]{0:T(128)S(1)} parameter(1) + %mul.1740 = f32[1,128,4096]{2,1,0:T(8,128)} broadcast(%param_1.1518), dimensions={1}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %mul.1739 = f32[1,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1021, %mul.1740), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %convert_element_type.1020 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1739), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %convert.149 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%convert_element_type.1020) + %dot_general.444 = f32[1,128,4096]{2,1,0:T(8,128)} multiply(%convert.148, %convert.149), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/...k,k->...k/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","4096"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %convert.150 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} convert(%dot_general.444) + ROOT %bitcast.794 = bf16[128,4096,1]{1,0,2:T(8,128)(2,1)} bitcast(%convert.150) +} + +%fused_computation.107.clone.clone.clone.1 (param_0.1437: bf16[1,4096,8,128]) -> bf16[4096,8,128] { + %param_0.1437 = bf16[1,4096,8,128]{1,3,2,0:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.793 = bf16[4096,8,128]{0,2,1:T(8,128)(2,1)} bitcast(%param_0.1437), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/dot_general" stack_frame_id=0} +} + +%fused_computation.208.clone.1 (param_0.1439: bf16[1,4096,8,128], param_1.1519: bf16[4096], param_2.1164: f32[128], param_3.786: bf16[4,1,128,4096], param_4.485: s32[]) -> bf16[1,128,8,128] { + %param_1.1519 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(1) + %param_2.1164 = f32[128]{0:T(128)S(1)} parameter(2) + %param_3.786 = bf16[4,1,128,4096]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(3) + %param_4.485 = s32[]{:T(128)S(6)} parameter(4) + %fusion.424 = bf16[128,4096,1]{1,0,2:T(8,128)(2,1)} fusion(%param_1.1519, %param_2.1164, %param_3.786, %param_4.485), kind=kLoop, calls=%fused_computation.193.clone.clone.clone.1 + %param_0.1439 = bf16[1,4096,8,128]{1,3,2,0:T(8,128)(2,1)S(1)} parameter(0) + %fusion.423 = bf16[4096,8,128]{0,2,1:T(8,128)(2,1)} fusion(%param_0.1439), kind=kLoop, calls=%fused_computation.107.clone.clone.clone.1, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/dot_general" stack_frame_id=0} + %convolution.143 = bf16[128,8,128]{2,0,1:T(8,128)(2,1)} convolution(%fusion.424, %fusion.423), window={size=8 pad=7_7 rhs_reversal=1}, dim_labels=bf0_i0o->b0f, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/dot_general" stack_frame_id=0} + ROOT %bitcast.796 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} bitcast(%convolution.143), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/dot_general" stack_frame_id=0} +} + +%fused_computation.221.clone.1 (param_0.1440: bf16[1,128,8,128]) -> (bf16[1,128,8,64], bf16[1,128,8,64]) { + %param_0.1440 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %slice.73 = bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)} slice(%param_0.1440), slice={[0:1], [0:128], [0:8], [64:128]}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/split" stack_frame_id=0} + %convert.211 = f32[1,128,8,64]{3,1,2,0:T(8,128)} convert(%slice.73) + %neg.131 = f32[1,128,8,64]{3,1,2,0:T(8,128)} negate(%convert.211), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/neg" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","8","64"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %convert.212 = bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} convert(%neg.131) + %slice.74 = bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} slice(%param_0.1440), slice={[0:1], [0:128], [0:8], [0:64]}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/split" stack_frame_id=0} + ROOT %tuple.185 = (bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)}, bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)}) tuple(%convert.212, %slice.74) +} + +%fused_computation.224.clone.1 (param_0.1441: bf16[1,128,8,64], param_1.1520: bf16[1,128,8,64], param_2.1165: bf16[1,128,8,128], param_3.787: bf16[128,128], param_4.486: bf16[128,128]) -> bf16[8,128,128] { + %param_2.1165 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(2) + %convert.213 = f32[1,128,8,128]{3,1,2,0:T(8,128)} convert(%param_2.1165) + %param_4.486 = bf16[128,128]{1,0:T(8,128)(2,1)S(1)} parameter(4) + %mul.1744 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_4.486), dimensions={1,3}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %convert.214 = f32[1,128,8,128]{3,1,2,0:T(8,128)} convert(%mul.1744) + %mul.1742 = f32[1,128,8,128]{3,1,2,0:T(8,128)} multiply(%convert.213, %convert.214), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","8","128"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %param_1.1520 = bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) + %constant.1000 = bf16[]{:T(256)} constant(-inf) + %pad.63 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_1.1520, %constant.1000), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/concatenate" stack_frame_id=0} + %convert.215 = f32[1,128,8,128]{3,1,2,0:T(8,128)} convert(%pad.63) + %param_0.1441 = bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %pad.62 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1441, %constant.1000), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/concatenate" stack_frame_id=0} + %convert.216 = f32[1,128,8,128]{3,1,2,0:T(8,128)} convert(%pad.62) + %maximum.50 = f32[1,128,8,128]{3,1,2,0:T(8,128)} maximum(%convert.215, %convert.216), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/concatenate" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","8","128"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %param_3.787 = bf16[128,128]{1,0:T(8,128)(2,1)S(1)} parameter(3) + %mul.1743 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_3.787), dimensions={1,3}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %convert.217 = f32[1,128,8,128]{3,1,2,0:T(8,128)} convert(%mul.1743) + %mul.1741 = f32[1,128,8,128]{3,1,2,0:T(8,128)} multiply(%maximum.50, %convert.217), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","8","128"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %add.887 = f32[1,128,8,128]{3,1,2,0:T(8,128)} add(%mul.1742, %mul.1741), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/add" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","8","128"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %convert.218 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)} convert(%add.887) + ROOT %bitcast.797 = bf16[8,128,128]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%convert.218), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/shard_map/vmap(static_sparsity_splash)/jit(_splash_attention)/splash_mha_fwd_segmented_residuals/splash_mha_fwd_segmented_residuals/squeeze" stack_frame_id=0} +} + +%fused_computation.116.clone.1 (param_0.1429: bf16[4,1024,8,128], param_1.1513: s32[]) -> bf16[1,1024,8,128] { + %param_0.1429 = bf16[4,1024,8,128]{3,1,2,0:T(8,128)(2,1)} parameter(0) + %param_1.1513 = s32[]{:T(128)S(6)} parameter(1) + %constant.362.clone.65 = s32[]{:T(128)} constant(0) + ROOT %dynamic-slice.168 = bf16[1,1024,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} dynamic-slice(%param_0.1429, %param_1.1513, %constant.362.clone.65, %constant.362.clone.65, %constant.362.clone.65), dynamic_slice_sizes={1,1024,8,128}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/jit(dynamic_index_in_dim)/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} +} + +%fused_computation.192.clone.clone.clone.1 (param_0.1431: bf16[4096], param_1.1514: f32[128], param_2.1161: bf16[4,1,128,4096], param_3.783: s32[]) -> bf16[128,4096,1] { + %param_0.1431 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(0) + %dot_general.443 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_0.1431), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/...k,k->...k/dot_general" stack_frame_id=0} + %convert.145 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%dot_general.443) + %param_2.1161 = bf16[4,1,128,4096]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %param_3.783 = s32[]{:T(128)S(6)} parameter(3) + %constant.362.clone.66 = s32[]{:T(128)} constant(0) + %dynamic-slice.169 = bf16[1,1,128,4096]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_2.1161, %param_3.783, %constant.362.clone.66, %constant.362.clone.66, %constant.362.clone.66), dynamic_slice_sizes={1,1,128,4096}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/jit(dynamic_index_in_dim)/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %bitcast.789 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic-slice.169), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/jit(dynamic_index_in_dim)/squeeze" stack_frame_id=0} + %convert_element_type.1017 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%bitcast.789), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %param_1.1514 = f32[128]{0:T(128)S(1)} parameter(1) + %mul.1737 = f32[1,128,4096]{2,1,0:T(8,128)} broadcast(%param_1.1514), dimensions={1}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %mul.1736 = f32[1,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1017, %mul.1737), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %convert_element_type.1016 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1736), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %convert.146 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%convert_element_type.1016) + %dot_general.442 = f32[1,128,4096]{2,1,0:T(8,128)} multiply(%convert.145, %convert.146), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/...k,k->...k/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","4096"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %convert.147 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} convert(%dot_general.442) + ROOT %bitcast.788 = bf16[128,4096,1]{1,0,2:T(8,128)(2,1)} bitcast(%convert.147) +} + +%fused_computation.103.clone.clone.clone.1 (param_0.1430: bf16[1,4096,8,128]) -> bf16[4096,8,128] { + %param_0.1430 = bf16[1,4096,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.787 = bf16[4096,8,128]{2,0,1:T(8,128)(2,1)} bitcast(%param_0.1430), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/dot_general" stack_frame_id=0} +} + +%fused_computation.209.clone.1 (param_0.1432: bf16[1,4096,8,128], param_1.1515: bf16[4096], param_2.1162: f32[128], param_3.784: bf16[4,1,128,4096], param_4.484: s32[]) -> bf16[8,128,128] { + %param_1.1515 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(1) + %param_2.1162 = f32[128]{0:T(128)S(1)} parameter(2) + %param_3.784 = bf16[4,1,128,4096]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(3) + %param_4.484 = s32[]{:T(128)S(6)} parameter(4) + %fusion.422 = bf16[128,4096,1]{1,0,2:T(8,128)(2,1)} fusion(%param_1.1515, %param_2.1162, %param_3.784, %param_4.484), kind=kLoop, calls=%fused_computation.192.clone.clone.clone.1 + %param_0.1432 = bf16[1,4096,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %fusion.421 = bf16[4096,8,128]{2,0,1:T(8,128)(2,1)} fusion(%param_0.1432), kind=kLoop, calls=%fused_computation.103.clone.clone.clone.1, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/dot_general" stack_frame_id=0} + %convolution.142 = bf16[128,8,128]{2,0,1:T(8,128)(2,1)} convolution(%fusion.422, %fusion.421), window={size=8 pad=7_7 rhs_reversal=1}, dim_labels=bf0_i0o->b0f, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/dot_general" stack_frame_id=0} + ROOT %bitcast.790 = bf16[8,128,128]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%convolution.142), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/shard_map/vmap(static_sparsity_splash)/jit(_splash_attention)/splash_mha_fwd_segmented_residuals/splash_mha_fwd_segmented_residuals/squeeze" stack_frame_id=0} +} + +%fused_computation.173.clone.clone.clone.1 (param_0.1450: bf16[32,128,128]) -> bf16[128,32,128] { + %param_0.1450 = bf16[32,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.804 = bf16[128,32,128]{2,0,1:T(8,128)(2,1)} bitcast(%param_0.1450) +} + +%fused_computation.90.clone.clone.clone.clone.clone.clone.clone.1 (param_0.1449: bf16[1,32,128,4096]) -> bf16[32,128,4096] { + %param_0.1449 = bf16[1,32,128,4096]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.803 = bf16[32,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%param_0.1449), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/dot_general" stack_frame_id=0} } %region_13.16 (reduce_sum.120: f32[], reduce_sum.121: f32[]) -> f32[] { - %reduce_sum.120 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/reduce_sum"} - %reduce_sum.121 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/reduce_sum"} - ROOT %reduce_sum.122 = f32[]{:T(128)} add(%reduce_sum.120, %reduce_sum.121), metadata={op_name="checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.25.clone.1.clone.clone.clone.clone.clone.clone (param_0.1263: bf16[4,32,128,4096], param_1.1391: s32[]) -> bf16[32,128,4096,1] { - %param_0.1263 = bf16[4,32,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1391 = s32[]{:T(128)S(6)} parameter(1) - %constant.1137 = s32[]{:T(128)} constant(0) - %dynamic_slice.325 = bf16[1,32,128,4096]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1263, %param_1.1391, %constant.1137, %constant.1137, %constant.1137), dynamic_slice_sizes={1,32,128,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.563 = bf16[32,128,4096,1]{2,1,0,3:T(8,128)(2,1)} bitcast(%dynamic_slice.325), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} -} - -%fused_computation.80.clone.clone.clone.clone.clone.clone (param_0.1264: bf16[4,32,128,128]) -> bf16[4,128,32,128] { - %param_0.1264 = bf16[4,32,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(0) - ROOT %bitcast.564 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} bitcast(%param_0.1264), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} -} - -%fused_computation.61.clone.clone (param_0.1265: bf16[4,32,128,4096], param_1.1392: s32[], param_2.1179: bf16[4,32,128,128], param_3.850: bf16[4,4,128,4096]) -> (f32[4,128], bf16[4,128,4096]) { - %param_3.850 = bf16[4,4,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(3) - %param_1.1392 = s32[]{:T(128)S(6)} parameter(1) - %constant.365.clone.1.clone.3 = s32[]{:T(128)} constant(0) - %dynamic_slice.208.clone.3 = bf16[1,4,128,4096]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_3.850, %param_1.1392, %constant.365.clone.1.clone.3, %constant.365.clone.1.clone.3, %constant.365.clone.1.clone.3), dynamic_slice_sizes={1,4,128,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - %bitcast.207.clone.3 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.208.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} - %param_2.1179 = bf16[4,32,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(2) - %fusion.83.clone.3 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} fusion(%param_2.1179), kind=kLoop, calls=%fused_computation.80.clone.clone.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} - %param_0.1265 = bf16[4,32,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(0) - %fusion.82.clone.3 = bf16[32,128,4096,1]{2,1,0,3:T(8,128)(2,1)} fusion(%param_0.1265, %param_1.1392), kind=kLoop, calls=%fused_computation.25.clone.1.clone.clone.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %convolution.62.clone.3 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} convolution(%fusion.83.clone.3, %fusion.82.clone.3), window={size=1x32}, dim_labels=0b1f_1io0->0bf1, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} - %bitcast.182.clone.3 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%convolution.62.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} - %add.635.clone.3 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} add(%bitcast.207.clone.3, %bitcast.182.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} - %convert_element_type.1102 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%add.635.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %square.215 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1102, %convert_element_type.1102), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/square" stack_frame_id=0} - %constant.1138 = f32[]{:T(128)} constant(0) - %reduce.177 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%square.215, %constant.1138), dimensions={2}, to_apply=%region_13.16, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0} - ROOT %tuple.188 = (f32[4,128]{1,0:T(4,128)S(1)}, bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)}) tuple(%reduce.177, %add.635.clone.3) -} - -%convert_element_type.523.reduce_sub_computation (lhs: bf16[], rhs: bf16[]) -> bf16[] { - %lhs = bf16[] parameter(0) - %rhs = bf16[] parameter(1) - ROOT %add.623 = bf16[] add(%lhs, %rhs), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.122.clone.clone (param_0.1247: bf16[4,4096], param_1.1379: s32[]) -> bf16[4096] { - %param_0.1247 = bf16[4,4096]{1,0:T(4,128)(2,1)} parameter(0) - %param_1.1379 = s32[]{:T(128)S(6)} parameter(1) - %constant.1126 = s32[]{:T(128)} constant(0) - %dynamic_slice.318 = bf16[1,4096]{1,0:T(2,128)(2,1)} dynamic-slice(%param_0.1247, %param_1.1379, %constant.1126), dynamic_slice_sizes={1,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - %constant.1127 = bf16[]{:T(256)} constant(-0), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %reduce.176 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} reduce(%dynamic_slice.318, %constant.1127), dimensions={0}, to_apply=%convert_element_type.523.reduce_sub_computation, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} -} - -%fused_computation.12.clone.clone.clone (param_0.1266: bf16[4,14336,4096], param_1.1393: s32[]) -> bf16[14336,4096,1] { - %param_0.1266 = bf16[4,14336,4096]{2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1393 = s32[]{:T(128)S(6)} parameter(1) - %constant.1139 = s32[]{:T(128)} constant(0) - %dynamic_slice.326 = bf16[1,14336,4096]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1266, %param_1.1393, %constant.1139, %constant.1139), dynamic_slice_sizes={1,14336,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.566 = bf16[14336,4096,1]{1,0,2:T(8,128)(2,1)} bitcast(%dynamic_slice.326), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} -} - -%bitcast_fusion.3.clone.clone (bitcast_input.12: bf16[4,128,4096]) -> bf16[4,128,4096] { - %bitcast_input.12 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) - ROOT %bitcast.565 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%bitcast_input.12) -} - -%fused_computation.13.clone.clone (param_0.1267: bf16[4,128,4096], param_1.1394: bf16[4,14336,4096], param_2.1180: s32[]) -> bf16[14336,4,128] { - %param_1.1394 = bf16[4,14336,4096]{2,1,0:T(8,128)(2,1)} parameter(1) - %param_2.1180 = s32[]{:T(128)S(6)} parameter(2) - %fusion.344 = bf16[14336,4096,1]{1,0,2:T(8,128)(2,1)} fusion(%param_1.1394, %param_2.1180), kind=kLoop, calls=%fused_computation.12.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_0.1267 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) - %fusion.345 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} fusion(%param_0.1267), kind=kLoop, calls=%bitcast_fusion.3.clone.clone - ROOT %convolution.114 = bf16[14336,4,128]{0,2,1:T(8,128)(2,1)S(1)} convolution(%fusion.344, %fusion.345), window={size=4 pad=3_3 rhs_reversal=1}, dim_labels=bf0_0oi->b0f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/dot_general" stack_frame_id=0} -} - -%fused_computation.144.clone.1.clone (param_0.1268: f32[4,128]) -> f32[4,128] { - %param_0.1268 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %constant.1141 = f32[]{:T(128)} constant(0.000244140625) - %closed_call.86 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1141), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} - %div.847 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.1268, %closed_call.86), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} - %constant.1140 = f32[]{:T(128)} constant(1e-05) - %closed_call.85 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1140), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} - %add.861 = f32[4,128]{1,0:T(4,128)} add(%div.847, %closed_call.85), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} - ROOT %rsqrt.98 = f32[4,128]{1,0:T(4,128)S(1)} rsqrt(%add.861), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/rsqrt" stack_frame_id=0} -} - -%fused_computation.11.clone.1.clone.clone (param_0.1272: bf16[4,4096,14336], param_1.1398: s32[]) -> bf16[4096,14336,1] { - %param_0.1272 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1398 = s32[]{:T(128)S(6)} parameter(1) - %constant.1143 = s32[]{:T(128)} constant(0) - %dynamic_slice.328 = bf16[1,4096,14336]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1272, %param_1.1398, %constant.1143, %constant.1143), dynamic_slice_sizes={1,4096,14336}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.568 = bf16[4096,14336,1]{1,0,2:T(8,128)(2,1)} bitcast(%dynamic_slice.328), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} -} - -%fused_computation.96.clone.2.clone.clone (param_0.1273: f32[4,128], param_1.1399: bf16[4,128,4096], param_2.1183: bf16[4096]) -> bf16[4,128,4096] { - %param_2.1183 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(2) - %dot_general.432 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1183), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1399 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %convert_element_type.1106 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_1.1399), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_0.1273 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.1717 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_0.1273), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.1716 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1106, %mul.1717), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1105 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1716), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %dot_general.431 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.432, %convert_element_type.1105), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} -} - -%fused_computation.23.clone.clone (param_0.1274: bf16[4,4096,14336], param_1.1400: s32[], param_2.1184: f32[4,128], param_3.852: bf16[4,128,4096], param_4.533: bf16[4096]) -> bf16[4,128,14336] { - %param_2.1184 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %param_3.852 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) - %param_4.533 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(4) - %fusion.349 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} fusion(%param_2.1184, %param_3.852, %param_4.533), kind=kLoop, calls=%fused_computation.96.clone.2.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_0.1274 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1400 = s32[]{:T(128)S(6)} parameter(1) - %fusion.348 = bf16[4096,14336,1]{1,0,2:T(8,128)(2,1)} fusion(%param_0.1274, %param_1.1400), kind=kLoop, calls=%fused_computation.11.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %convolution.116 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)S(1)} convolution(%fusion.349, %fusion.348), window={size=1}, dim_labels=0bf_io0->0bf, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} -} - -%fused_computation.14.clone.1.clone.clone (param_0.1275: bf16[4,4096,14336], param_1.1401: s32[]) -> bf16[4096,14336,1] { - %param_0.1275 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1401 = s32[]{:T(128)S(6)} parameter(1) - %constant.1144 = s32[]{:T(128)} constant(0) - %dynamic_slice.329 = bf16[1,4096,14336]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1275, %param_1.1401, %constant.1144, %constant.1144), dynamic_slice_sizes={1,4096,14336}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.569 = bf16[4096,14336,1]{1,0,2:T(8,128)(2,1)} bitcast(%dynamic_slice.329), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} -} - -%fused_computation.39.clone.1.clone.clone (param_0.1276: bf16[14336,4,128], param_1.1402: bf16[4,128,14336]) -> bf16[4,128,14336] { - %param_1.1402 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %constant.1145 = bf16[]{:T(256)} constant(1) - %jit_silu_.44 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} broadcast(%constant.1145), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)" stack_frame_id=0} - %neg.130 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} negate(%param_1.1402), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/neg" stack_frame_id=0} - %exp.69 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} exponential(%neg.130), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/exp" stack_frame_id=0} - %add.862 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} add(%exp.69, %jit_silu_.44), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/add" stack_frame_id=0} - %div.848 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} divide(%jit_silu_.44, %add.862), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/div" stack_frame_id=0} - %mul.1719 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} multiply(%param_1.1402, %div.848), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/mul" stack_frame_id=0} - %param_0.1276 = bf16[14336,4,128]{0,2,1:T(8,128)(2,1)S(1)} parameter(0) - %bitcast.570 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} bitcast(%param_0.1276), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/dot_general" stack_frame_id=0} - ROOT %mul.1718 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} multiply(%mul.1719, %bitcast.570), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} -} - -%fused_computation.21.clone.clone (param_0.1277: bf16[4,4096,14336], param_1.1403: s32[], param_2.1185: bf16[14336,4,128], param_3.853: bf16[4,128,14336]) -> bf16[4,128,4096] { - %param_2.1185 = bf16[14336,4,128]{0,2,1:T(8,128)(2,1)S(1)} parameter(2) - %param_3.853 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) - %bitcast_multiply_fusion.15 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} fusion(%param_2.1185, %param_3.853), kind=kLoop, calls=%fused_computation.39.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} - %param_0.1277 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1403 = s32[]{:T(128)S(6)} parameter(1) - %fusion.350 = bf16[4096,14336,1]{1,0,2:T(8,128)(2,1)} fusion(%param_0.1277, %param_1.1403), kind=kLoop, calls=%fused_computation.14.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %convolution.117 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} convolution(%bitcast_multiply_fusion.15, %fusion.350), window={size=1}, dim_labels=0bf_oi0->0bf, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/dot_general" stack_frame_id=0} -} - -%fused_computation.14.clone.clone.clone (param_0.1269: bf16[4,4096,14336], param_1.1395: s32[]) -> bf16[4096,14336,1] { - %param_0.1269 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1395 = s32[]{:T(128)S(6)} parameter(1) - %constant.1142 = s32[]{:T(128)} constant(0) - %dynamic_slice.327 = bf16[1,4096,14336]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1269, %param_1.1395, %constant.1142, %constant.1142), dynamic_slice_sizes={1,4096,14336}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.567 = bf16[4096,14336,1]{1,0,2:T(8,128)(2,1)} bitcast(%dynamic_slice.327), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} -} - -%fused_computation.96.clone.1.clone.clone (param_0.1270: f32[4,128], param_1.1396: bf16[4,128,4096], param_2.1181: bf16[4096]) -> bf16[4,128,4096] { - %param_2.1181 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(2) - %dot_general.430 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1181), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1396 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %convert_element_type.1104 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_1.1396), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_0.1270 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.1715 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_0.1270), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.1714 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1104, %mul.1715), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1103 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1714), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %dot_general.429 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.430, %convert_element_type.1103), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} -} - -%fused_computation.20.clone.clone (param_0.1271: bf16[4,4096,14336], param_1.1397: s32[], param_2.1182: f32[4,128], param_3.851: bf16[4,128,4096], param_4.532: bf16[4096]) -> bf16[4,128,14336] { - %param_2.1182 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %param_3.851 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) - %param_4.532 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(4) - %fusion.347 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} fusion(%param_2.1182, %param_3.851, %param_4.532), kind=kLoop, calls=%fused_computation.96.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_0.1271 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1397 = s32[]{:T(128)S(6)} parameter(1) - %fusion.346 = bf16[4096,14336,1]{1,0,2:T(8,128)(2,1)} fusion(%param_0.1271, %param_1.1397), kind=kLoop, calls=%fused_computation.14.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %convolution.115 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)S(1)} convolution(%fusion.347, %fusion.346), window={size=1}, dim_labels=0bf_io0->0bf, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} + %reduce_sum.120 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/reduce_sum"} + %reduce_sum.121 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/reduce_sum"} + ROOT %reduce_sum.122 = f32[]{:T(128)} add(%reduce_sum.120, %reduce_sum.121), metadata={op_name="checkpoint/rematted_computation/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%fused_computation.156.clone.1 (param_0.1451: bf16[1,32,128,4096], param_1.1526: bf16[32,128,128], param_2.1169: bf16[4,1,128,4096], param_3.791: s32[]) -> (f32[128], bf16[1,128,4096]) { + %param_2.1169 = bf16[4,1,128,4096]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %param_3.791 = s32[]{:T(128)S(6)} parameter(3) + %constant.362.clone.16.clone.3 = s32[]{:T(128)} constant(0) + %dynamic-slice.114.clone.3 = bf16[1,1,128,4096]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_2.1169, %param_3.791, %constant.362.clone.16.clone.3, %constant.362.clone.16.clone.3, %constant.362.clone.16.clone.3), dynamic_slice_sizes={1,1,128,4096}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/jit(dynamic_index_in_dim)/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %bitcast.391.clone.3 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic-slice.114.clone.3), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/jit(dynamic_index_in_dim)/squeeze" stack_frame_id=0} + %convert.219 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%bitcast.391.clone.3) + %param_1.1526 = bf16[32,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %fusion.190.clone.3 = bf16[128,32,128]{2,0,1:T(8,128)(2,1)} fusion(%param_1.1526), kind=kLoop, calls=%fused_computation.173.clone.clone.clone.1 + %param_0.1451 = bf16[1,32,128,4096]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %fusion.170.clone.3 = bf16[32,128,4096]{2,1,0:T(8,128)(2,1)} fusion(%param_0.1451), kind=kLoop, calls=%fused_computation.90.clone.clone.clone.clone.clone.clone.clone.1, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/dot_general" stack_frame_id=0} + %convolution.105.clone.3 = bf16[128,4096,1]{1,0,2:T(8,128)(2,1)} convolution(%fusion.190.clone.3, %fusion.170.clone.3), window={size=32}, dim_labels=b0f_0io->bf0, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/dot_general" stack_frame_id=0} + %bitcast.341.clone.3 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%convolution.105.clone.3), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/dot_general" stack_frame_id=0} + %convert.220 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%bitcast.341.clone.3) + %add.685.clone.3 = f32[1,128,4096]{2,1,0:T(8,128)} add(%convert.219, %convert.220), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/add" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","4096"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %convert_element_type.1024 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%add.685.clone.3), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %square.216 = f32[1,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1024, %convert_element_type.1024), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/square" stack_frame_id=0} + %constant.363.clone.12 = f32[]{:T(128)} constant(0) + %reduce.190 = f32[128]{0:T(128)S(1)} reduce(%square.216, %constant.363.clone.12), dimensions={0,2}, to_apply=%region_13.16, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/reduce_sum" stack_frame_id=0} + %convert.221 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)S(1)} convert(%add.685.clone.3) + ROOT %tuple.187 = (f32[128]{0:T(128)S(1)}, bf16[1,128,4096]{2,1,0:T(8,128)(2,1)S(1)}) tuple(%reduce.190, %convert.221) +} + +%fused_computation.241.clone.1 (param_0.1452: f32[128]) -> f32[128] { + %param_0.1452 = f32[128]{0:T(128)S(1)} parameter(0) + %constant.364.clone.10 = f32[]{:T(128)} constant(0.000244140625) + %broadcast.711 = f32[128]{0:T(128)} broadcast(%constant.364.clone.10), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call" stack_frame_id=0} + %div.857 = f32[128]{0:T(128)} multiply(%param_0.1452, %broadcast.711), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/div" stack_frame_id=0} + %constant.365.clone.10 = f32[]{:T(128)} constant(1e-05) + %broadcast.710 = f32[128]{0:T(128)} broadcast(%constant.365.clone.10), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call" stack_frame_id=0} + %add.889 = f32[128]{0:T(128)} add(%div.857, %broadcast.710), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/add" stack_frame_id=0} + %bitcast.806 = f32[1,128]{1,0:T(1,128)} bitcast(%add.889), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/add" stack_frame_id=0} + %rsqrt.115 = f32[1,128]{1,0:T(1,128)} rsqrt(%bitcast.806), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/rsqrt" stack_frame_id=0} + ROOT %bitcast.805 = f32[128]{0:T(128)S(1)} bitcast(%rsqrt.115), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/rsqrt" stack_frame_id=0} +} + +%fused_computation.184.clone.1 (param_0.1455: bf16[4096], param_1.1528: f32[128], param_2.1170: bf16[1,128,4096]) -> bf16[128,4096] { + %param_0.1455 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(0) + %dot_general.449 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_0.1455), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/...k,k->...k/dot_general" stack_frame_id=0} + %convert.154 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%dot_general.449) + %param_2.1170 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %convert_element_type.1026 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%param_2.1170), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %param_1.1528 = f32[128]{0:T(128)S(1)} parameter(1) + %mul.1752 = f32[1,128,4096]{2,1,0:T(8,128)} broadcast(%param_1.1528), dimensions={1}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %mul.1751 = f32[1,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1026, %mul.1752), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %convert_element_type.1025 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1751), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %convert.155 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%convert_element_type.1025) + %dot_general.448 = f32[1,128,4096]{2,1,0:T(8,128)} multiply(%convert.154, %convert.155), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/...k,k->...k/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","4096"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %convert.156 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} convert(%dot_general.448) + ROOT %bitcast.808 = bf16[128,4096]{1,0:T(8,128)(2,1)} bitcast(%convert.156) +} + +%fused_computation.72.clone.clone.clone.1 (param_0.1454: bf16[1,4096,14336]) -> bf16[4096,14336] { + %param_0.1454 = bf16[1,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) + ROOT %bitcast.807 = bf16[4096,14336]{1,0:T(8,128)(2,1)} bitcast(%param_0.1454), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/dot_general" stack_frame_id=0} +} + +%fused_computation.110.clone.1 (param_0.1456: bf16[1,4096,14336], param_1.1529: bf16[4096], param_2.1171: bf16[1,128,4096], param_3.792: f32[128]) -> bf16[1,128,14336] { + %param_1.1529 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(1) + %param_3.792 = f32[128]{0:T(128)S(1)} parameter(3) + %param_2.1171 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %fusion.428 = bf16[128,4096]{1,0:T(8,128)(2,1)} fusion(%param_1.1529, %param_3.792, %param_2.1171), kind=kLoop, calls=%fused_computation.184.clone.1 + %param_0.1456 = bf16[1,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) + %fusion.427 = bf16[4096,14336]{1,0:T(8,128)(2,1)} fusion(%param_0.1456), kind=kLoop, calls=%fused_computation.72.clone.clone.clone.1, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/dot_general" stack_frame_id=0} + %convolution.145 = bf16[128,14336]{1,0:T(8,128)(2,1)} convolution(%fusion.428, %fusion.427), dim_labels=bf_io->bf, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/dot_general" stack_frame_id=0} + ROOT %bitcast.809 = bf16[1,128,14336]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%convolution.145), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/dot_general" stack_frame_id=0} +} + +%fused_computation.86.clone.1 (param_0.1461: bf16[4,1024,14336], param_1.1532: s32[]) -> bf16[1,1024,14336] { + %param_0.1461 = bf16[4,1024,14336]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1532 = s32[]{:T(128)S(6)} parameter(1) + %constant.362.clone.74 = s32[]{:T(128)} constant(0) + ROOT %dynamic-slice.177 = bf16[1,1024,14336]{2,1,0:T(8,128)(2,1)S(1)} dynamic-slice(%param_0.1461, %param_1.1532, %constant.362.clone.74, %constant.362.clone.74), dynamic_slice_sizes={1,1024,14336}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/jit(dynamic_index_in_dim)/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} +} + +%fused_computation.185.clone.1 (param_0.1463: bf16[4096], param_1.1533: f32[128], param_2.1172: bf16[1,128,4096]) -> bf16[128,4096] { + %param_0.1463 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(0) + %dot_general.451 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_0.1463), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/...k,k->...k/dot_general" stack_frame_id=0} + %convert.157 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%dot_general.451) + %param_2.1172 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %convert_element_type.1028 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%param_2.1172), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %param_1.1533 = f32[128]{0:T(128)S(1)} parameter(1) + %mul.1754 = f32[1,128,4096]{2,1,0:T(8,128)} broadcast(%param_1.1533), dimensions={1}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %mul.1753 = f32[1,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1028, %mul.1754), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %convert_element_type.1027 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1753), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %convert.158 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%convert_element_type.1027) + %dot_general.450 = f32[1,128,4096]{2,1,0:T(8,128)} multiply(%convert.157, %convert.158), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/...k,k->...k/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","4096"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %convert.159 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} convert(%dot_general.450) + ROOT %bitcast.814 = bf16[128,4096]{1,0:T(8,128)(2,1)} bitcast(%convert.159) +} + +%fused_computation.80.clone.clone.clone.1 (param_0.1462: bf16[1,4096,14336]) -> bf16[4096,14336] { + %param_0.1462 = bf16[1,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) + ROOT %bitcast.813 = bf16[4096,14336]{1,0:T(8,128)(2,1)} bitcast(%param_0.1462), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/dot_general" stack_frame_id=0} +} + +%fused_computation.108.clone.1 (param_0.1464: bf16[1,4096,14336], param_1.1534: bf16[4096], param_2.1173: bf16[1,128,4096], param_3.793: f32[128]) -> bf16[1,128,14336] { + %param_1.1534 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(1) + %param_3.793 = f32[128]{0:T(128)S(1)} parameter(3) + %param_2.1173 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %fusion.432 = bf16[128,4096]{1,0:T(8,128)(2,1)} fusion(%param_1.1534, %param_3.793, %param_2.1173), kind=kLoop, calls=%fused_computation.185.clone.1 + %param_0.1464 = bf16[1,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) + %fusion.431 = bf16[4096,14336]{1,0:T(8,128)(2,1)} fusion(%param_0.1464), kind=kLoop, calls=%fused_computation.80.clone.clone.clone.1, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/dot_general" stack_frame_id=0} + %convolution.147 = bf16[128,14336]{1,0:T(8,128)(2,1)} convolution(%fusion.432, %fusion.431), dim_labels=bf_io->bf, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/dot_general" stack_frame_id=0} + ROOT %bitcast.815 = bf16[1,128,14336]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%convolution.147), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/dot_general" stack_frame_id=0} +} + +%fused_computation.122.clone.1 (param_0.1509: bf16[1,128,14336], param_1.1561: bf16[1,128,14336]) -> bf16[128,14336] { + %param_1.1561 = bf16[1,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %convert.178 = f32[1,128,14336]{2,1,0:T(8,128)} convert(%param_1.1561) + %constant.371.clone.23 = bf16[]{:T(256)} constant(1) + %jit_silu_.44 = bf16[1,128,14336]{2,1,0:T(8,128)(2,1)} broadcast(%constant.371.clone.23), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/jit(silu)" stack_frame_id=0} + %convert.179 = f32[1,128,14336]{2,1,0:T(8,128)} convert(%jit_silu_.44) + %neg.135 = f32[1,128,14336]{2,1,0:T(8,128)} negate(%convert.178), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/jit(silu)/neg" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","14336"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %exp.81 = f32[1,128,14336]{2,1,0:T(8,128)} exponential(%neg.135), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/jit(silu)/exp" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","14336"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %add.894 = f32[1,128,14336]{2,1,0:T(8,128)} add(%exp.81, %convert.179), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/jit(silu)/add" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","14336"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %div.864 = f32[1,128,14336]{2,1,0:T(8,128)} divide(%convert.179, %add.894), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/jit(silu)/div" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","14336"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %mul.1787 = f32[1,128,14336]{2,1,0:T(8,128)} multiply(%convert.178, %div.864), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/jit(silu)/mul" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","14336"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %param_0.1509 = bf16[1,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert.180 = f32[1,128,14336]{2,1,0:T(8,128)} convert(%param_0.1509) + %mul.1786 = f32[1,128,14336]{2,1,0:T(8,128)} multiply(%mul.1787, %convert.180), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","14336"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %convert.181 = bf16[1,128,14336]{2,1,0:T(8,128)(2,1)} convert(%mul.1786) + ROOT %bitcast.867 = bf16[128,14336]{1,0:T(8,128)(2,1)} bitcast(%convert.181) +} + +%fused_computation.179.clone.1 (param_0.1508: bf16[1,128,4096]) -> bf16[128,4096] { + %param_0.1508 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.866 = bf16[128,4096]{1,0:T(8,128)(2,1)} bitcast(%param_0.1508) +} + +%fused_computation.121.clone.1 (param_0.1510: bf16[1,128,14336], param_1.1562: bf16[1,128,4096], param_2.1192: bf16[1,128,14336]) -> bf16[14336,4096] { + %param_0.1510 = bf16[1,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %param_2.1192 = bf16[1,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %fusion.447 = bf16[128,14336]{1,0:T(8,128)(2,1)} fusion(%param_0.1510, %param_2.1192), kind=kLoop, calls=%fused_computation.122.clone.1 + %param_1.1562 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %fusion.448 = bf16[128,4096]{1,0:T(8,128)(2,1)} fusion(%param_1.1562), kind=kLoop, calls=%fused_computation.179.clone.1 + ROOT %convolution.155 = bf16[14336,4096]{1,0:T(8,128)(2,1)} convolution(%fusion.447, %fusion.448), dim_labels=fb_io->bf, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/dot_general" stack_frame_id=0} +} + +%fused_computation.253.clone.1 (param_0.1476: f32[1,32,128,128]) -> (f32[32,128], f32[1,32,128,1]) { + %param_0.1476 = f32[1,32,128,128]{2,1,3,0:T(8,128)S(1)} parameter(0) + %slice.77 = f32[1,32,128,1]{2,1,3,0:T(8,128)} slice(%param_0.1476), slice={[0:1], [0:32], [0:128], [0:1]}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/shard_map/vmap(static_sparsity_splash)/jit(_splash_attention)/slice" stack_frame_id=0} + %bitcast.827 = f32[32,128]{1,0:T(8,128)} bitcast(%slice.77), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/shard_map/vmap(static_sparsity_splash)/jit(_splash_attention)/broadcast_in_dim" stack_frame_id=0} + ROOT %tuple.191 = (f32[32,128]{1,0:T(8,128)}, f32[1,32,128,1]{2,1,3,0:T(8,128)}) tuple(%bitcast.827, %slice.77) +} + +%fused_computation.85.clone.1 (param_0.1457: bf16[4,14336,1024], param_1.1530: s32[]) -> bf16[1,14336,1024] { + %param_0.1457 = bf16[4,14336,1024]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1530 = s32[]{:T(128)S(6)} parameter(1) + %constant.362.clone.73 = s32[]{:T(128)} constant(0) + ROOT %dynamic-slice.176 = bf16[1,14336,1024]{2,1,0:T(8,128)(2,1)S(1)} dynamic-slice(%param_0.1457, %param_1.1530, %constant.362.clone.73, %constant.362.clone.73), dynamic_slice_sizes={1,14336,1024}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/jit(dynamic_index_in_dim)/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} +} + +%fused_computation.76.clone.clone.clone.1 (param_0.1459: bf16[1,14336,4096]) -> bf16[14336,4096] { + %param_0.1459 = bf16[1,14336,4096]{2,1,0:T(8,128)(2,1)} parameter(0) + ROOT %bitcast.811 = bf16[14336,4096]{1,0:T(8,128)(2,1)} bitcast(%param_0.1459), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/dot_general" stack_frame_id=0} +} + +%fused_computation.180.clone.1 (param_0.1458: bf16[1,128,4096]) -> bf16[128,4096] { + %param_0.1458 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.810 = bf16[128,4096]{1,0:T(8,128)(2,1)} bitcast(%param_0.1458) +} + +%fused_computation.109.clone.1 (param_0.1460: bf16[1,14336,4096], param_1.1531: bf16[1,128,4096]) -> bf16[1,128,14336] { + %param_0.1460 = bf16[1,14336,4096]{2,1,0:T(8,128)(2,1)} parameter(0) + %fusion.429 = bf16[14336,4096]{1,0:T(8,128)(2,1)} fusion(%param_0.1460), kind=kLoop, calls=%fused_computation.76.clone.clone.clone.1, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/dot_general" stack_frame_id=0} + %param_1.1531 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %fusion.430 = bf16[128,4096]{1,0:T(8,128)(2,1)} fusion(%param_1.1531), kind=kLoop, calls=%fused_computation.180.clone.1 + %convolution.146 = bf16[14336,128]{0,1:T(8,128)(2,1)} convolution(%fusion.429, %fusion.430), dim_labels=bf_oi->bf, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/dot_general" stack_frame_id=0} + ROOT %bitcast.812 = bf16[1,128,14336]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%convolution.146), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/dot_general" stack_frame_id=0} +} + +%fused_computation.125.clone.1 (param_0.1466: bf16[1,128,14336], param_1.1535: bf16[1,128,14336]) -> bf16[128,14336] { + %param_1.1535 = bf16[1,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %convert.160 = f32[1,128,14336]{2,1,0:T(8,128)} convert(%param_1.1535) + %constant.371.clone.21 = bf16[]{:T(256)} constant(1) + %jit_silu_.42 = bf16[1,128,14336]{2,1,0:T(8,128)(2,1)} broadcast(%constant.371.clone.21), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/jit(silu)" stack_frame_id=0} + %convert.161 = f32[1,128,14336]{2,1,0:T(8,128)} convert(%jit_silu_.42) + %neg.133 = f32[1,128,14336]{2,1,0:T(8,128)} negate(%convert.160), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/jit(silu)/neg" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","14336"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %exp.79 = f32[1,128,14336]{2,1,0:T(8,128)} exponential(%neg.133), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/jit(silu)/exp" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","14336"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %add.890 = f32[1,128,14336]{2,1,0:T(8,128)} add(%exp.79, %convert.161), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/jit(silu)/add" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","14336"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %div.858 = f32[1,128,14336]{2,1,0:T(8,128)} divide(%convert.161, %add.890), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/jit(silu)/div" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","14336"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %mul.1756 = f32[1,128,14336]{2,1,0:T(8,128)} multiply(%convert.160, %div.858), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/jit(silu)/mul" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","14336"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %param_0.1466 = bf16[1,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert.162 = f32[1,128,14336]{2,1,0:T(8,128)} convert(%param_0.1466) + %mul.1755 = f32[1,128,14336]{2,1,0:T(8,128)} multiply(%mul.1756, %convert.162), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/mul" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","14336"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %convert.163 = bf16[1,128,14336]{2,1,0:T(8,128)(2,1)} convert(%mul.1755) + ROOT %bitcast.817 = bf16[128,14336]{1,0:T(8,128)(2,1)} bitcast(%convert.163) +} + +%fused_computation.74.clone.1 (param_0.1465: bf16[1,4096,14336]) -> bf16[4096,14336] { + %param_0.1465 = bf16[1,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) + ROOT %bitcast.816 = bf16[4096,14336]{1,0:T(8,128)(2,1)} bitcast(%param_0.1465), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/dot_general" stack_frame_id=0} +} + +%fused_computation.73.clone.1 (param_0.1467: bf16[1,4096,14336], param_1.1536: bf16[1,128,14336], param_2.1174: bf16[1,128,14336]) -> bf16[128,4096] { + %param_1.1536 = bf16[1,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %param_2.1174 = bf16[1,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %fusion.434 = bf16[128,14336]{1,0:T(8,128)(2,1)} fusion(%param_1.1536, %param_2.1174), kind=kLoop, calls=%fused_computation.125.clone.1 + %param_0.1467 = bf16[1,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) + %fusion.433 = bf16[4096,14336]{1,0:T(8,128)(2,1)} fusion(%param_0.1467), kind=kLoop, calls=%fused_computation.74.clone.1, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/dot_general" stack_frame_id=0} + ROOT %convolution.148 = bf16[128,4096]{1,0:T(8,128)(2,1)S(1)} convolution(%fusion.434, %fusion.433), dim_labels=bf_oi->bf, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/dot_general" stack_frame_id=0} +} + +%fused_computation.120.clone.1.clone.clone.clone.1 (param_0.1469: bf16[1,128,14336], param_1.1537: bf16[1,128,14336], param_2.1175: bf16[1,128,14336]) -> bf16[128,14336] { + %param_1.1537 = bf16[1,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %convert.164 = f32[1,128,14336]{2,1,0:T(8,128)} convert(%param_1.1537) + %param_2.1175 = bf16[1,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %convert.165 = f32[1,128,14336]{2,1,0:T(8,128)} convert(%param_2.1175) + %mul.1761 = f32[1,128,14336]{2,1,0:T(8,128)} multiply(%convert.164, %convert.165), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/mul" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","14336"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %constant.371.clone.22 = bf16[]{:T(256)} constant(1) + %jit_silu_.43 = bf16[1,128,14336]{2,1,0:T(8,128)(2,1)} broadcast(%constant.371.clone.22), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/jit(silu)" stack_frame_id=0} + %convert.167 = f32[1,128,14336]{2,1,0:T(8,128)} convert(%jit_silu_.43) + %param_0.1469 = bf16[1,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert.166 = f32[1,128,14336]{2,1,0:T(8,128)} convert(%param_0.1469) + %neg.134 = f32[1,128,14336]{2,1,0:T(8,128)} negate(%convert.166), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/jit(silu)/neg" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","14336"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %exp.80 = f32[1,128,14336]{2,1,0:T(8,128)} exponential(%neg.134), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/jit(silu)/exp" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","14336"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %add.891 = f32[1,128,14336]{2,1,0:T(8,128)} add(%exp.80, %convert.167), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/jit(silu)/add" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","14336"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %div.859 = f32[1,128,14336]{2,1,0:T(8,128)} divide(%convert.167, %add.891), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/jit(silu)/div" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","14336"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %mul.1760 = f32[1,128,14336]{2,1,0:T(8,128)} multiply(%mul.1761, %div.859), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/jit(silu)/mul" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","14336"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %mul.1759 = f32[1,128,14336]{2,1,0:T(8,128)} multiply(%convert.166, %mul.1761), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/jit(silu)/mul" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","14336"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %sub.118 = f32[1,128,14336]{2,1,0:T(8,128)} subtract(%convert.167, %div.859), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/jit(silu)/sub" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","14336"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %mul.1758 = f32[1,128,14336]{2,1,0:T(8,128)} multiply(%div.859, %sub.118), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/jit(silu)/mul" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","14336"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %mul.1757 = f32[1,128,14336]{2,1,0:T(8,128)} multiply(%mul.1759, %mul.1758), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/jit(silu)/mul" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","14336"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %add_any.160 = f32[1,128,14336]{2,1,0:T(8,128)} add(%mul.1760, %mul.1757), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/jit(silu)/add_any" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","14336"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %convert.168 = bf16[1,128,14336]{2,1,0:T(8,128)(2,1)} convert(%add_any.160) + ROOT %bitcast.819 = bf16[128,14336]{1,0:T(8,128)(2,1)} bitcast(%convert.168) +} + +%fused_computation.78.clone.clone.clone.clone.clone.clone.clone.1 (param_0.1468: bf16[1,4096,14336]) -> bf16[4096,14336] { + %param_0.1468 = bf16[1,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) + ROOT %bitcast.818 = bf16[4096,14336]{1,0:T(8,128)(2,1)} bitcast(%param_0.1468), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/dot_general" stack_frame_id=0} } %region_14.17 (reduce_sum.126: f32[], reduce_sum.127: f32[]) -> f32[] { - %reduce_sum.126 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/reduce_sum"} - %reduce_sum.127 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/reduce_sum"} - ROOT %reduce_sum.128 = f32[]{:T(128)} add(%reduce_sum.126, %reduce_sum.127), metadata={op_name="checkpoint/layers/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.11.clone.clone.clone.clone.clone.clone.clone (param_0.1278: bf16[4,4096,14336], param_1.1404: s32[]) -> bf16[4096,14336,1] { - %param_0.1278 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1404 = s32[]{:T(128)S(6)} parameter(1) - %constant.1146 = s32[]{:T(128)} constant(0) - %dynamic_slice.330 = bf16[1,4096,14336]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1278, %param_1.1404, %constant.1146, %constant.1146), dynamic_slice_sizes={1,4096,14336}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.571 = bf16[4096,14336,1]{1,0,2:T(8,128)(2,1)} bitcast(%dynamic_slice.330), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} -} - -%fused_computation.38.clone.1.clone.clone.clone.clone (param_0.1279: bf16[4,128,14336], param_1.1405: bf16[4,128,14336], param_2.1186: bf16[14336,4,128]) -> bf16[4,128,14336] { - %param_2.1186 = bf16[14336,4,128]{0,2,1:T(8,128)(2,1)S(1)} parameter(2) - %bitcast.572 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} bitcast(%param_2.1186), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/dot_general" stack_frame_id=0} - %param_1.1405 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %mul.1724 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} multiply(%bitcast.572, %param_1.1405), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} - %constant.1147 = bf16[]{:T(256)} constant(1) - %jit_silu_.45 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} broadcast(%constant.1147), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)" stack_frame_id=0} - %param_0.1279 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) - %neg.131 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} negate(%param_0.1279), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/neg" stack_frame_id=0} - %exp.70 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} exponential(%neg.131), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/exp" stack_frame_id=0} - %add.863 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} add(%exp.70, %jit_silu_.45), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/add" stack_frame_id=0} - %div.849 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} divide(%jit_silu_.45, %add.863), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/div" stack_frame_id=0} - %mul.1723 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} multiply(%mul.1724, %div.849), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/jit(silu)/mul" stack_frame_id=0} - %mul.1722 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} multiply(%param_0.1279, %mul.1724), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/jit(silu)/mul" stack_frame_id=0} - %sub.98 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} subtract(%jit_silu_.45, %div.849), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/sub" stack_frame_id=0} - %mul.1721 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} multiply(%div.849, %sub.98), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/mul" stack_frame_id=0} - %mul.1720 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} multiply(%mul.1722, %mul.1721), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/jit(silu)/mul" stack_frame_id=0} - ROOT %add_any.145 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} add(%mul.1723, %mul.1720), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/jit(silu)/add_any" stack_frame_id=0} -} - -%fused_computation.63.clone.clone (param_0.1280: bf16[4,128,4096], param_1.1406: bf16[4096], param_2.1187: bf16[4,128,4096], param_3.854: bf16[4,4096,14336], param_4.534: s32[], param_5.435: bf16[4,128,14336], param_6.304: bf16[4,128,14336], param_7.200: bf16[14336,4,128]) -> (f32[4,128], bf16[4,128,4096]) { - %param_0.1280 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) - %convert_element_type.1108 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_0.1280), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_2.1187 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) - %param_5.435 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(5) - %param_6.304 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(6) - %param_7.200 = bf16[14336,4,128]{0,2,1:T(8,128)(2,1)S(1)} parameter(7) - %fusion.134.clone.3 = bf16[4,128,14336]{2,1,0:T(8,128)(2,1)} fusion(%param_5.435, %param_6.304, %param_7.200), kind=kLoop, calls=%fused_computation.38.clone.1.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/jit(silu)/add_any" stack_frame_id=0} - %param_3.854 = bf16[4,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(3) - %param_4.534 = s32[]{:T(128)S(6)} parameter(4) - %fusion.79.clone.3 = bf16[4096,14336,1]{1,0,2:T(8,128)(2,1)} fusion(%param_3.854, %param_4.534), kind=kLoop, calls=%fused_computation.11.clone.clone.clone.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %convolution.60.clone.3 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convolution(%fusion.134.clone.3, %fusion.79.clone.3), window={size=1}, dim_labels=0bf_oi0->0bf, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/dot_general" stack_frame_id=0} - %add_any.132.clone.3 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} add(%param_2.1187, %convolution.60.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/add_any" stack_frame_id=0} - %param_1.1406 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(1) - %dot_general.434 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_1.1406), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %dot_general.433 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%add_any.132.clone.3, %dot_general.434), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/...k,k->...k/dot_general" stack_frame_id=0} - %convert_element_type.1107 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%dot_general.433), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/convert_element_type" stack_frame_id=0} - %mul.1725 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1108, %convert_element_type.1107), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} - %constant.1148 = f32[]{:T(128)} constant(0) - %reduce.178 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%mul.1725, %constant.1148), dimensions={2}, to_apply=%region_14.17, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/reduce_sum" stack_frame_id=0} - ROOT %tuple.189 = (f32[4,128]{1,0:T(4,128)S(1)}, bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)}) tuple(%reduce.178, %add_any.132.clone.3) -} - -%fused_computation.140.clone.clone (param_0.1281: f32[4,128], param_1.1407: f32[4,128]) -> f32[4,128] { - %param_0.1281 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %param_1.1407 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %constant.1152 = f32[]{:T(128)} constant(0.000244140625) - %closed_call.89 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1152), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} - %div.851 = f32[4,128]{1,0:T(4,128)} multiply(%param_1.1407, %closed_call.89), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} - %constant.1151 = f32[]{:T(128)} constant(1e-05) - %closed_call.88 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1151), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} - %add.864 = f32[4,128]{1,0:T(4,128)} add(%div.851, %closed_call.88), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} - %rsqrt.99 = f32[4,128]{1,0:T(4,128)} rsqrt(%add.864), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/rsqrt" stack_frame_id=0} - %div.850 = f32[4,128]{1,0:T(4,128)} divide(%rsqrt.99, %add.864), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} - %constant.1150 = f32[]{:T(128)} constant(-0.5) - %closed_call.87 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1150), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} - %mul.1728 = f32[4,128]{1,0:T(4,128)} multiply(%div.850, %closed_call.87), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.1727 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.1281, %mul.1728), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} - %constant.1149 = f32[]{:T(128)} constant(0.00048828125) - %mul.1729 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1149), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} - ROOT %mul.1726 = f32[4,128]{1,0:T(4,128)S(1)} multiply(%mul.1727, %mul.1729), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} -} - -%region_20.24 (dot_general.187: bf16[], dot_general.188: bf16[]) -> bf16[] { - %dot_general.187 = bf16[]{:T(256)} parameter(0), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/...k,k->...k/dot_general"} - %dot_general.188 = bf16[]{:T(256)} parameter(1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/...k,k->...k/dot_general"} - ROOT %add.173 = bf16[]{:T(256)} add(%dot_general.187, %dot_general.188), metadata={op_name="add.39"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.94.clone.clone (param_0.1282: bf16[4,128,4096], param_1.1408: f32[4,128], param_2.1188: bf16[4,128,4096], param_3.855: bf16[4,128,4096], param_4.535: f32[4,128], param_5.436: bf16[4096]) -> (bf16[4096], bf16[4,128,4096]) { - %param_0.1282 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) - %param_2.1188 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) - %convert_element_type.1110 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%param_2.1188), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_1.1408 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %mul.1731 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_1.1408), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.1730 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1110, %mul.1731), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1109 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1730), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %multiply.271 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%param_0.1282, %convert_element_type.1109), metadata={op_name="multiply.204"} - %constant.1153 = bf16[]{:T(256)} constant(0) - %reduce.179 = bf16[4096]{0:T(1024)(128)(2,1)} reduce(%multiply.271, %constant.1153), dimensions={0,1}, to_apply=%region_20.24, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_3.855 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) - %param_5.436 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(5) - %dot_general.286.clone.3 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_5.436), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %dot_general.263.clone.3 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} multiply(%param_0.1282, %dot_general.286.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/...k,k->...k/dot_general" stack_frame_id=0} - %convert_element_type.753.clone.3 = f32[4,128,4096]{2,1,0:T(8,128)} convert(%dot_general.263.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/convert_element_type" stack_frame_id=0} - %mul.1142.clone.3 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.753.clone.3, %mul.1731), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} - %param_4.535 = f32[4,128]{1,0:T(4,128)S(1)} parameter(4) - %mul.1151.clone.3 = f32[4,128,4096]{2,1,0:T(8,128)} broadcast(%param_4.535), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} - %mul.1141.clone.3 = f32[4,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1110, %mul.1151.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/mul" stack_frame_id=0} - %add_any.126.clone.3 = f32[4,128,4096]{2,1,0:T(8,128)} add(%mul.1142.clone.3, %mul.1141.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/add_any" stack_frame_id=0} - %convert_element_type.751.clone.3 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)} convert(%add_any.126.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/convert_element_type" stack_frame_id=0} - %add_any.124.clone.3 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} add(%param_3.855, %convert_element_type.751.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/add_any" stack_frame_id=0} - ROOT %tuple.190 = (bf16[4096]{0:T(1024)(128)(2,1)}, bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)}) tuple(%reduce.179, %add_any.124.clone.3) + %reduce_sum.126 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/reduce_sum"} + %reduce_sum.127 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/reduce_sum"} + ROOT %reduce_sum.128 = f32[]{:T(128)} add(%reduce_sum.126, %reduce_sum.127), metadata={op_name="checkpoint/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%fused_computation.158.clone.1 (param_0.1470: bf16[1,128,4096], param_1.1538: bf16[4096], param_2.1176: bf16[128,4096], param_3.794: bf16[1,4096,14336], param_4.489: bf16[1,128,14336], param_5.416: bf16[1,128,14336], param_6.268: bf16[1,128,14336]) -> (f32[128], bf16[1,128,4096]) { + %param_0.1470 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.1030 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%param_0.1470), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %param_2.1176 = bf16[128,4096]{1,0:T(8,128)(2,1)S(1)} parameter(2) + %convert.222 = f32[128,4096]{1,0:T(8,128)} convert(%param_2.1176) + %param_4.489 = bf16[1,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(4) + %param_5.416 = bf16[1,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(5) + %param_6.268 = bf16[1,128,14336]{2,1,0:T(8,128)(2,1)S(1)} parameter(6) + %fusion.271.clone.3 = bf16[128,14336]{1,0:T(8,128)(2,1)} fusion(%param_4.489, %param_5.416, %param_6.268), kind=kLoop, calls=%fused_computation.120.clone.1.clone.clone.clone.1 + %param_3.794 = bf16[1,4096,14336]{2,1,0:T(8,128)(2,1)} parameter(3) + %fusion.158.clone.3 = bf16[4096,14336]{1,0:T(8,128)(2,1)} fusion(%param_3.794), kind=kLoop, calls=%fused_computation.78.clone.clone.clone.clone.clone.clone.clone.1, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/dot_general" stack_frame_id=0} + %convolution.97.clone.3 = bf16[128,4096]{1,0:T(8,128)(2,1)} convolution(%fusion.271.clone.3, %fusion.158.clone.3), dim_labels=bf_oi->bf, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/dot_general" stack_frame_id=0} + %convert.223 = f32[128,4096]{1,0:T(8,128)} convert(%convolution.97.clone.3) + %add_any.122.clone.3 = f32[128,4096]{1,0:T(8,128)} add(%convert.222, %convert.223), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/add_any" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["128","4096"],"tuple_shapes":[],"layout":{"minor_to_major":["1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false]}}} + %convert.224 = bf16[128,4096]{1,0:T(8,128)(2,1)} convert(%add_any.122.clone.3) + %bitcast.323.clone.3 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%convert.224), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/add_any" stack_frame_id=0} + %convert.225 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%bitcast.323.clone.3) + %param_1.1538 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(1) + %dot_general.453 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_1.1538), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/...k,k->...k/dot_general" stack_frame_id=0} + %convert.226 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%dot_general.453) + %dot_general.452 = f32[1,128,4096]{2,1,0:T(8,128)} multiply(%convert.225, %convert.226), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/...k,k->...k/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","4096"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %convert_element_type.1029 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%dot_general.452), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/convert_element_type" stack_frame_id=0} + %mul.1762 = f32[1,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1030, %convert_element_type.1029), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/mul" stack_frame_id=0} + %constant.363.clone.13 = f32[]{:T(128)} constant(0) + %reduce.191 = f32[128]{0:T(128)S(1)} reduce(%mul.1762, %constant.363.clone.13), dimensions={0,2}, to_apply=%region_14.17, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/reduce_sum" stack_frame_id=0} + ROOT %tuple.188 = (f32[128]{0:T(128)S(1)}, bf16[1,128,4096]{2,1,0:T(8,128)(2,1)S(1)}) tuple(%reduce.191, %bitcast.323.clone.3) +} + +%fused_computation.251.clone.1 (param_0.1471: f32[128], param_1.1539: f32[128]) -> f32[128] { + %param_0.1471 = f32[128]{0:T(128)S(1)} parameter(0) + %bitcast.822 = f32[1,128]{1,0:T(1,128)} bitcast(%param_0.1471), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/reduce_sum" stack_frame_id=0} + %param_1.1539 = f32[128]{0:T(128)S(1)} parameter(1) + %constant.364.clone.11 = f32[]{:T(128)} constant(0.000244140625) + %broadcast.713 = f32[128]{0:T(128)} broadcast(%constant.364.clone.11), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call" stack_frame_id=0} + %div.861 = f32[128]{0:T(128)} multiply(%param_1.1539, %broadcast.713), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/div" stack_frame_id=0} + %constant.365.clone.11 = f32[]{:T(128)} constant(1e-05) + %broadcast.712 = f32[128]{0:T(128)} broadcast(%constant.365.clone.11), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call" stack_frame_id=0} + %add.892 = f32[128]{0:T(128)} add(%div.861, %broadcast.712), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/add" stack_frame_id=0} + %bitcast.821 = f32[1,128]{1,0:T(1,128)} bitcast(%add.892), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/add" stack_frame_id=0} + %rsqrt.116 = f32[1,128]{1,0:T(1,128)} rsqrt(%bitcast.821), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/rsqrt" stack_frame_id=0} + %div.860 = f32[1,128]{1,0:T(1,128)} divide(%rsqrt.116, %bitcast.821), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/div" stack_frame_id=0} + %constant.372.clone.5 = f32[]{:T(128)} constant(-0.5) + %closed_call.65 = f32[1,128]{1,0:T(1,128)} broadcast(%constant.372.clone.5), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call" stack_frame_id=0} + %mul.1765 = f32[1,128]{1,0:T(1,128)} multiply(%div.860, %closed_call.65), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %mul.1764 = f32[1,128]{1,0:T(1,128)} multiply(%bitcast.822, %mul.1765), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/mul" stack_frame_id=0} + %constant.373.clone.5 = f32[]{:T(128)} constant(0.00048828125) + %mul.1766 = f32[1,128]{1,0:T(1,128)} broadcast(%constant.373.clone.5), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/mul" stack_frame_id=0} + %mul.1763 = f32[1,128]{1,0:T(1,128)} multiply(%mul.1764, %mul.1766), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/mul" stack_frame_id=0} + ROOT %bitcast.820 = f32[128]{0:T(128)S(1)} bitcast(%mul.1763), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/mul" stack_frame_id=0} +} + +%region_17.21 (dot_general.187: bf16[], dot_general.188: bf16[]) -> bf16[] { + %dot_general.187 = bf16[]{:T(256)} parameter(0), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/...k,k->...k/dot_general"} + %dot_general.188 = bf16[]{:T(256)} parameter(1), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/...k,k->...k/dot_general"} + ROOT %add.167 = bf16[]{:T(256)} add(%dot_general.187, %dot_general.188), metadata={op_name="add.44"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%fused_computation.196.clone.1 (param_0.1472: bf16[1,128,4096], param_1.1540: f32[128], param_2.1177: bf16[1,128,4096], param_3.795: bf16[1,128,4096], param_4.490: f32[128], param_5.417: bf16[4096]) -> (bf16[4096], bf16[1,128,4096]) { + %param_0.1472 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert.227 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%param_0.1472) + %param_2.1177 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %convert_element_type.1032 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%param_2.1177), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %param_1.1540 = f32[128]{0:T(128)S(1)} parameter(1) + %mul.1768 = f32[1,128,4096]{2,1,0:T(8,128)} broadcast(%param_1.1540), dimensions={1}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %mul.1767 = f32[1,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1032, %mul.1768), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %convert_element_type.1031 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} convert(%mul.1767), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %convert.228 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%convert_element_type.1031) + %multiply.298 = f32[1,128,4096]{2,1,0:T(8,128)} multiply(%convert.227, %convert.228), metadata={op_name="multiply.195"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","4096"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %convert.229 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} convert(%multiply.298) + %constant.374.clone.5 = bf16[]{:T(256)} constant(0) + %reduce.192 = bf16[4096]{0:T(1024)(128)(2,1)} reduce(%convert.229, %constant.374.clone.5), dimensions={0,1}, to_apply=%region_17.21, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/...k,k->...k/dot_general" stack_frame_id=0} + %param_3.795 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) + %convert.231 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%param_3.795) + %param_5.417 = bf16[4096]{0:T(1024)(128)(2,1)S(1)} parameter(5) + %dot_general.378.clone.3 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} broadcast(%param_5.417), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/...k,k->...k/dot_general" stack_frame_id=0} + %convert.230 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%dot_general.378.clone.3) + %dot_general.294.clone.3 = f32[1,128,4096]{2,1,0:T(8,128)} multiply(%convert.227, %convert.230), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/...k,k->...k/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","4096"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %convert_element_type.766.clone.3 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%dot_general.294.clone.3), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/convert_element_type" stack_frame_id=0} + %mul.1195.clone.3 = f32[1,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.766.clone.3, %mul.1768), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/mul" stack_frame_id=0} + %param_4.490 = f32[128]{0:T(128)S(1)} parameter(4) + %mul.1208.clone.3 = f32[1,128,4096]{2,1,0:T(8,128)} broadcast(%param_4.490), dimensions={1}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/mul" stack_frame_id=0} + %mul.1194.clone.3 = f32[1,128,4096]{2,1,0:T(8,128)} multiply(%convert_element_type.1032, %mul.1208.clone.3), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/mul" stack_frame_id=0} + %add_any.142.clone.3 = f32[1,128,4096]{2,1,0:T(8,128)} add(%mul.1195.clone.3, %mul.1194.clone.3), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/add_any" stack_frame_id=0} + %convert_element_type.764.clone.3 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} convert(%add_any.142.clone.3), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/convert_element_type" stack_frame_id=0} + %convert.232 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%convert_element_type.764.clone.3) + %add_any.140.clone.3 = f32[1,128,4096]{2,1,0:T(8,128)} add(%convert.231, %convert.232), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/add_any" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","4096"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %convert.233 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} convert(%add_any.140.clone.3) + ROOT %tuple.189 = (bf16[4096]{0:T(1024)(128)(2,1)}, bf16[1,128,4096]{2,1,0:T(8,128)(2,1)}) tuple(%reduce.192, %convert.233) +} + +%fused_computation.178.clone.clone.clone.1 (param_0.1474: bf16[1,128,4096]) -> bf16[128,4096,1] { + %param_0.1474 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} parameter(0) + ROOT %bitcast.824 = bf16[128,4096,1]{1,0,2:T(8,128)(2,1)} bitcast(%param_0.1474) +} + +%fused_computation.88.clone.clone.clone.1 (param_0.1473: bf16[1,32,128,4096]) -> bf16[32,128,4096] { + %param_0.1473 = bf16[1,32,128,4096]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.823 = bf16[32,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%param_0.1473), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/dot_general" stack_frame_id=0} } %region_15.18 (dot_general.184: f32[], dot_general.185: f32[]) -> f32[] { - %dot_general.184 = f32[]{:T(128)} parameter(0), metadata={op_name="vmap(jit(_splash_attention))/hsd,hsd->hs/dot_general"} - %dot_general.185 = f32[]{:T(128)} parameter(1), metadata={op_name="vmap(jit(_splash_attention))/hsd,hsd->hs/dot_general"} - ROOT %add.169 = f32[]{:T(128)} add(%dot_general.184, %dot_general.185), metadata={op_name="add.31"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.25.clone.clone.clone.clone.clone.clone.clone (param_0.1283: bf16[4,32,128,4096], param_1.1409: s32[]) -> bf16[32,128,4096,1] { - %param_0.1283 = bf16[4,32,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1409 = s32[]{:T(128)S(6)} parameter(1) - %constant.1154 = s32[]{:T(128)} constant(0) - %dynamic_slice.331 = bf16[1,32,128,4096]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1283, %param_1.1409, %constant.1154, %constant.1154, %constant.1154), dynamic_slice_sizes={1,32,128,4096}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.573 = bf16[32,128,4096,1]{2,1,0,3:T(8,128)(2,1)} bitcast(%dynamic_slice.331), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} -} - -%fused_computation.76.clone.clone.clone.clone.clone.clone (param_0.1284: bf16[4,128,4096]) -> bf16[4,128,4096,1] { - %param_0.1284 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) - ROOT %bitcast.574 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} bitcast(%param_0.1284), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/add_any" stack_frame_id=0} -} - -%fused_computation.66.clone.clone (param_0.1285: bf16[4,32,128,128], param_1.1410: bf16[4,32,128,4096], param_2.1189: s32[], param_3.856: bf16[4,128,4096]) -> (f32[4,32,128], bf16[4,32,128,128]) { - %param_0.1285 = bf16[4,32,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(0) - %convert.124 = f32[4,32,128,128]{3,2,1,0:T(8,128)} convert(%param_0.1285), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/shard_map/convert" stack_frame_id=0} - %param_3.856 = bf16[4,128,4096]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) - %fusion.95.clone.3 = bf16[4,128,4096,1]{2,1,3,0:T(8,128)(2,1)} fusion(%param_3.856), kind=kLoop, calls=%fused_computation.76.clone.clone.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/add_any" stack_frame_id=0} - %param_1.1410 = bf16[4,32,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(1) - %param_2.1189 = s32[]{:T(128)S(6)} parameter(2) - %fusion.94.clone.3 = bf16[32,128,4096,1]{2,1,0,3:T(8,128)(2,1)} fusion(%param_1.1410, %param_2.1189), kind=kLoop, calls=%fused_computation.25.clone.clone.clone.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %convolution.64.clone.3 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} convolution(%fusion.95.clone.3, %fusion.94.clone.3), window={size=1x32 pad=0_0x31_31 rhs_reversal=0x1}, dim_labels=0bf1_1oi0->0b1f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/dot_general" stack_frame_id=0} - %constant.619.clone.3 = bf16[]{:T(256)} constant(0.25) - %div.442.clone.3 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%constant.619.clone.3), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/div" stack_frame_id=0} - %div.441.clone.3 = bf16[4,128,32,128]{3,1,2,0:T(8,128)(2,1)} multiply(%convolution.64.clone.3, %div.442.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/div" stack_frame_id=0} - %bitcast.209.clone.3 = bf16[4,32,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} bitcast(%div.441.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/div" stack_frame_id=0} - %convert.123 = f32[4,32,128,128]{3,2,1,0:T(8,128)} convert(%bitcast.209.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/shard_map/convert.1" stack_frame_id=0} - %multiply.272 = f32[4,32,128,128]{3,2,1,0:T(8,128)} multiply(%convert.124, %convert.123), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/shard_map/multiply" stack_frame_id=0} - %constant.1155 = f32[]{:T(128)} constant(0) - %dot_general.435 = f32[4,32,128]{2,1,0:T(8,128)S(1)} reduce(%multiply.272, %constant.1155), dimensions={3}, to_apply=%region_15.18, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/shard_map/vmap(jit(_splash_attention))/hsd,hsd->hs/dot_general" stack_frame_id=0} - ROOT %tuple.191 = (f32[4,32,128]{2,1,0:T(8,128)S(1)}, bf16[4,32,128,128]{3,2,1,0:T(8,128)(2,1)S(1)}) tuple(%dot_general.435, %bitcast.209.clone.3) -} + %dot_general.184 = f32[]{:T(128)} parameter(0), metadata={op_name="vmap(static_sparsity_splash)/jit(_splash_attention)/hsd,hsd->hs/dot_general"} + %dot_general.185 = f32[]{:T(128)} parameter(1), metadata={op_name="vmap(static_sparsity_splash)/jit(_splash_attention)/hsd,hsd->hs/dot_general"} + ROOT %add.166 = f32[]{:T(128)} add(%dot_general.184, %dot_general.185), metadata={op_name="add.31"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%fused_computation.159.clone.1 (param_0.1475: bf16[32,128,128], param_1.1541: bf16[1,32,128,4096], param_2.1178: bf16[1,128,4096]) -> (f32[32,128], bf16[128,32,128]) { + %param_0.1475 = bf16[32,128,128]{2,1,0:T(8,128)(2,1)} parameter(0) + %bitcast.826 = bf16[1,32,128,128]{3,2,1,0:T(8,128)(2,1)} bitcast(%param_0.1475), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/shard_map/vmap(static_sparsity_splash)/jit(_splash_attention)/splash_mha_fwd_segmented_residuals/splash_mha_fwd_segmented_residuals/broadcast_in_dim" stack_frame_id=0} + %convert.125 = f32[1,32,128,128]{3,2,1,0:T(8,128)} convert(%bitcast.826), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/shard_map/convert" stack_frame_id=0} + %param_2.1178 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} parameter(2) + %fusion.193.clone.3 = bf16[128,4096,1]{1,0,2:T(8,128)(2,1)} fusion(%param_2.1178), kind=kLoop, calls=%fused_computation.178.clone.clone.clone.1 + %param_1.1541 = bf16[1,32,128,4096]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %fusion.107.clone.3 = bf16[32,128,4096]{2,1,0:T(8,128)(2,1)} fusion(%param_1.1541), kind=kLoop, calls=%fused_computation.88.clone.clone.clone.1, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/dot_general" stack_frame_id=0} + %convolution.70.clone.3 = bf16[128,32,128]{2,0,1:T(8,128)(2,1)} convolution(%fusion.193.clone.3, %fusion.107.clone.3), window={size=32 pad=31_31 rhs_reversal=1}, dim_labels=bf0_0oi->b0f, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/dot_general" stack_frame_id=0} + %bitcast.825 = bf16[1,32,128,128]{3,2,1,0:T(8,128)(2,1)} bitcast(%convolution.70.clone.3), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/transpose" stack_frame_id=0} + %convert.124 = f32[1,32,128,128]{3,2,1,0:T(8,128)} convert(%bitcast.825), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/shard_map/convert.1" stack_frame_id=0} + %multiply.299 = f32[1,32,128,128]{3,2,1,0:T(8,128)} multiply(%convert.125, %convert.124), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/shard_map/multiply" stack_frame_id=0} + %constant.363.clone.14 = f32[]{:T(128)} constant(0) + %dot_general.454 = f32[32,128]{1,0:T(8,128)S(1)} reduce(%multiply.299, %constant.363.clone.14), dimensions={0,3}, to_apply=%region_15.18, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/shard_map/vmap(static_sparsity_splash)/jit(_splash_attention)/hsd,hsd->hs/dot_general" stack_frame_id=0} + ROOT %tuple.190 = (f32[32,128]{1,0:T(8,128)S(1)}, bf16[128,32,128]{2,0,1:T(8,128)(2,1)}) tuple(%dot_general.454, %convolution.70.clone.3) +} + +%fused_computation.176.clone.1 (param_0.1488: bf16[32,128,128], param_1.1549: bf16[128,128]) -> bf16[1,128,32,128] { + %param_0.1488 = bf16[32,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %bitcast.842 = bf16[1,32,128,128]{3,2,1,0:T(8,128)(2,1)} bitcast(%param_0.1488), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/shard_map/vmap(static_sparsity_splash)/jit(_splash_attention)/splash_mha_dq_segmented_no_residuals/splash_mha_dq_segmented_no_residuals/broadcast_in_dim" stack_frame_id=0} + %convert.234 = f32[1,32,128,128]{3,2,1,0:T(8,128)} convert(%bitcast.842) + %param_1.1549 = bf16[128,128]{1,0:T(8,128)(2,1)S(1)} parameter(1) + %broadcast.716 = bf16[1,32,128,128]{3,2,1,0:T(8,128)(2,1)} broadcast(%param_1.1549), dimensions={2,3}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %convert.235 = f32[1,32,128,128]{3,2,1,0:T(8,128)} convert(%broadcast.716) + %mul.1775 = f32[1,32,128,128]{3,2,1,0:T(8,128)} multiply(%convert.234, %convert.235), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/mul" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","32","128","128"],"tuple_shapes":[],"layout":{"minor_to_major":["3","2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %convert.236 = bf16[1,32,128,128]{3,2,1,0:T(8,128)(2,1)} convert(%mul.1775) + ROOT %bitcast.841 = bf16[1,128,32,128]{3,1,2,0:T(8,128)(2,1)S(1)} bitcast(%convert.236), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/mul" stack_frame_id=0} +} + +%fused_computation.418.clone.1 (param_0.1489: bf16[1,128,32,128]) -> (bf16[1,128,32,64], bf16[1,128,32,64]) { + %param_0.1489 = bf16[1,128,32,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %slice.79 = bf16[1,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)} slice(%param_0.1489), slice={[0:1], [0:128], [0:32], [64:128]}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/split" stack_frame_id=0} + %slice.50.clone.3 = bf16[1,128,32,64]{3,1,2,0:T(8,128)(2,1)} slice(%param_0.1489), slice={[0:1], [0:128], [0:32], [0:64]}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/split" stack_frame_id=0} + %convert.237 = f32[1,128,32,64]{3,1,2,0:T(8,128)} convert(%slice.50.clone.3) + %neg.103.clone.3 = f32[1,128,32,64]{3,1,2,0:T(8,128)} negate(%convert.237), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/neg" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","32","64"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %convert.238 = bf16[1,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)} convert(%neg.103.clone.3) + ROOT %tuple.195 = (bf16[1,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)}, bf16[1,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)}) tuple(%slice.79, %convert.238) +} + +%fused_computation.170.clone.1 (param_0.1490: bf16[1,128,32,64], param_1.1550: bf16[1,128,32,64], param_2.1185: bf16[32,128,128], param_3.802: bf16[128,128]) -> bf16[128,32,128] { + %param_1.1550 = bf16[1,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) + %constant.1005 = bf16[]{:T(256)} constant(-inf) + %pad.69 = bf16[1,128,32,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_1.1550, %constant.1005), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/concatenate" stack_frame_id=0} + %convert.239 = f32[1,128,32,128]{3,1,2,0:T(8,128)} convert(%pad.69) + %param_0.1490 = bf16[1,128,32,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %pad.68 = bf16[1,128,32,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1490, %constant.1005), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/concatenate" stack_frame_id=0} + %convert.240 = f32[1,128,32,128]{3,1,2,0:T(8,128)} convert(%pad.68) + %maximum.53 = f32[1,128,32,128]{3,1,2,0:T(8,128)} maximum(%convert.239, %convert.240), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/concatenate" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","32","128"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %param_2.1185 = bf16[32,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %bitcast.845 = bf16[1,32,128,128]{3,2,1,0:T(8,128)(2,1)} bitcast(%param_2.1185), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/shard_map/vmap(static_sparsity_splash)/jit(_splash_attention)/splash_mha_dq_segmented_no_residuals/splash_mha_dq_segmented_no_residuals/broadcast_in_dim" stack_frame_id=0} + %convert.241 = f32[1,32,128,128]{3,2,1,0:T(8,128)} convert(%bitcast.845) + %param_3.802 = bf16[128,128]{1,0:T(8,128)(2,1)S(1)} parameter(3) + %broadcast.717 = bf16[1,32,128,128]{3,2,1,0:T(8,128)(2,1)} broadcast(%param_3.802), dimensions={2,3}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %convert.242 = f32[1,32,128,128]{3,2,1,0:T(8,128)} convert(%broadcast.717) + %mul.1776 = f32[1,32,128,128]{3,2,1,0:T(8,128)} multiply(%convert.241, %convert.242), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/mul" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","32","128","128"],"tuple_shapes":[],"layout":{"minor_to_major":["3","2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %convert.243 = bf16[1,32,128,128]{3,2,1,0:T(8,128)(2,1)} convert(%mul.1776) + %bitcast.844 = bf16[1,128,32,128]{3,1,2,0:T(8,128)(2,1)} bitcast(%convert.243), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/mul" stack_frame_id=0} + %convert.244 = f32[1,128,32,128]{3,1,2,0:T(8,128)} convert(%bitcast.844) + %add_any.162 = f32[1,128,32,128]{3,1,2,0:T(8,128)} add(%maximum.53, %convert.244), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/add_any" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","32","128"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %convert.245 = bf16[1,128,32,128]{3,1,2,0:T(8,128)(2,1)} convert(%add_any.162) + ROOT %bitcast.843 = bf16[128,32,128]{2,0,1:T(8,128)(2,1)S(1)} bitcast(%convert.245) +} + +%bitcast_fusion.1.clone.1 (bitcast_input.13: bf16[128,32,128]) -> bf16[128,32,128] { + %bitcast_input.13 = bf16[128,32,128]{2,0,1:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.857 = bf16[128,32,128]{2,0,1:T(8,128)(2,1)} bitcast(%bitcast_input.13) +} + +%fused_computation.92.clone.clone.clone.1 (param_0.1501: bf16[1,4096,32,128]) -> bf16[4096,32,128] { + %param_0.1501 = bf16[1,4096,32,128]{1,3,2,0:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.856 = bf16[4096,32,128]{0,2,1:T(8,128)(2,1)} bitcast(%param_0.1501), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/dot_general" stack_frame_id=0} +} + +%fused_computation.141.clone.1 (param_0.1502: bf16[128,32,128], param_1.1557: bf16[1,4096,32,128]) -> bf16[128,4096] { + %param_0.1502 = bf16[128,32,128]{2,0,1:T(8,128)(2,1)S(1)} parameter(0) + %fusion.446 = bf16[128,32,128]{2,0,1:T(8,128)(2,1)} fusion(%param_0.1502), kind=kLoop, calls=%bitcast_fusion.1.clone.1 + %param_1.1557 = bf16[1,4096,32,128]{1,3,2,0:T(8,128)(2,1)S(1)} parameter(1) + %fusion.445 = bf16[4096,32,128]{0,2,1:T(8,128)(2,1)} fusion(%param_1.1557), kind=kLoop, calls=%fused_computation.92.clone.clone.clone.1, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/dot_general" stack_frame_id=0} + %convolution.154 = bf16[128,4096,1]{1,0,2:T(8,128)(2,1)} convolution(%fusion.446, %fusion.445), window={size=32}, dim_labels=b0f_o0i->bf0, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/dot_general" stack_frame_id=0} + ROOT %bitcast.858 = bf16[128,4096]{1,0:T(8,128)(2,1)} bitcast(%convolution.154), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/dot_general" stack_frame_id=0} +} + +%fused_computation.231.clone.1 (param_0.1480: bf16[8,128,128], param_1.1544: bf16[128,128]) -> bf16[1,128,8,128] { + %param_0.1480 = bf16[8,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %bitcast.832 = bf16[1,8,128,128]{3,2,1,0:T(8,128)(2,1)} bitcast(%param_0.1480), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/shard_map/vmap(static_sparsity_splash)/jit(_splash_attention)/splash_mha_dkv_segmented_no_residuals/splash_mha_dkv_segmented_no_residuals/broadcast_in_dim" stack_frame_id=0} + %convert.246 = f32[1,8,128,128]{3,2,1,0:T(8,128)} convert(%bitcast.832) + %param_1.1544 = bf16[128,128]{1,0:T(8,128)(2,1)S(1)} parameter(1) + %broadcast.714 = bf16[1,8,128,128]{3,2,1,0:T(8,128)(2,1)} broadcast(%param_1.1544), dimensions={2,3}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %convert.247 = f32[1,8,128,128]{3,2,1,0:T(8,128)} convert(%broadcast.714) + %mul.1771 = f32[1,8,128,128]{3,2,1,0:T(8,128)} multiply(%convert.246, %convert.247), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/mul" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","8","128","128"],"tuple_shapes":[],"layout":{"minor_to_major":["3","2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %convert.248 = bf16[1,8,128,128]{3,2,1,0:T(8,128)(2,1)} convert(%mul.1771) + ROOT %bitcast.831 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} bitcast(%convert.248), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/mul" stack_frame_id=0} +} + +%fused_computation.419.clone.1 (param_0.1481: bf16[1,128,8,128]) -> (bf16[1,128,8,64], bf16[1,128,8,64]) { + %param_0.1481 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %slice.78 = bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} slice(%param_0.1481), slice={[0:1], [0:128], [0:8], [64:128]}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/split" stack_frame_id=0} + %slice.52.clone.3 = bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)} slice(%param_0.1481), slice={[0:1], [0:128], [0:8], [0:64]}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/split" stack_frame_id=0} + %convert.249 = f32[1,128,8,64]{3,1,2,0:T(8,128)} convert(%slice.52.clone.3) + %neg.105.clone.3 = f32[1,128,8,64]{3,1,2,0:T(8,128)} negate(%convert.249), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/neg" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","8","64"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %convert.250 = bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} convert(%neg.105.clone.3) + ROOT %tuple.192 = (bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)}, bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)}) tuple(%slice.78, %convert.250) +} + +%fused_computation.226.clone.1 (param_0.1482: bf16[1,128,8,64], param_1.1545: bf16[1,128,8,64], param_2.1181: bf16[8,128,128], param_3.798: bf16[128,128]) -> bf16[128,8,128] { + %param_1.1545 = bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) + %constant.1002 = bf16[]{:T(256)} constant(-inf) + %pad.67 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_1.1545, %constant.1002), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/concatenate" stack_frame_id=0} + %convert.251 = f32[1,128,8,128]{3,1,2,0:T(8,128)} convert(%pad.67) + %param_0.1482 = bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %pad.66 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1482, %constant.1002), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/concatenate" stack_frame_id=0} + %convert.252 = f32[1,128,8,128]{3,1,2,0:T(8,128)} convert(%pad.66) + %maximum.52 = f32[1,128,8,128]{3,1,2,0:T(8,128)} maximum(%convert.251, %convert.252), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/concatenate" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","8","128"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %param_2.1181 = bf16[8,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %bitcast.835 = bf16[1,8,128,128]{3,2,1,0:T(8,128)(2,1)} bitcast(%param_2.1181), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/shard_map/vmap(static_sparsity_splash)/jit(_splash_attention)/splash_mha_dkv_segmented_no_residuals/splash_mha_dkv_segmented_no_residuals/broadcast_in_dim" stack_frame_id=0} + %convert.253 = f32[1,8,128,128]{3,2,1,0:T(8,128)} convert(%bitcast.835) + %param_3.798 = bf16[128,128]{1,0:T(8,128)(2,1)S(1)} parameter(3) + %broadcast.715 = bf16[1,8,128,128]{3,2,1,0:T(8,128)(2,1)} broadcast(%param_3.798), dimensions={2,3}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %convert.254 = f32[1,8,128,128]{3,2,1,0:T(8,128)} convert(%broadcast.715) + %mul.1772 = f32[1,8,128,128]{3,2,1,0:T(8,128)} multiply(%convert.253, %convert.254), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/mul" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","8","128","128"],"tuple_shapes":[],"layout":{"minor_to_major":["3","2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %convert.255 = bf16[1,8,128,128]{3,2,1,0:T(8,128)(2,1)} convert(%mul.1772) + %bitcast.834 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)} bitcast(%convert.255), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/mul" stack_frame_id=0} + %convert.256 = f32[1,128,8,128]{3,1,2,0:T(8,128)} convert(%bitcast.834) + %add_any.161 = f32[1,128,8,128]{3,1,2,0:T(8,128)} add(%maximum.52, %convert.256), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/add_any" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","8","128"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %convert.257 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)} convert(%add_any.161) + ROOT %bitcast.833 = bf16[128,8,128]{2,0,1:T(8,128)(2,1)S(1)} bitcast(%convert.257) +} + +%bitcast_fusion.clone.1 (bitcast_input.12: bf16[128,8,128]) -> bf16[128,8,128] { + %bitcast_input.12 = bf16[128,8,128]{2,0,1:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.854 = bf16[128,8,128]{2,0,1:T(8,128)(2,1)} bitcast(%bitcast_input.12) +} + +%fused_computation.105.clone.clone.clone.1 (param_0.1499: bf16[1,4096,8,128]) -> bf16[4096,8,128] { + %param_0.1499 = bf16[1,4096,8,128]{1,3,2,0:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.853 = bf16[4096,8,128]{0,2,1:T(8,128)(2,1)} bitcast(%param_0.1499), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/dot_general" stack_frame_id=0} +} + +%fused_computation.139.clone.1 (param_0.1500: bf16[128,8,128], param_1.1556: bf16[1,4096,8,128]) -> bf16[128,4096] { + %param_0.1500 = bf16[128,8,128]{2,0,1:T(8,128)(2,1)S(1)} parameter(0) + %fusion.444 = bf16[128,8,128]{2,0,1:T(8,128)(2,1)} fusion(%param_0.1500), kind=kLoop, calls=%bitcast_fusion.clone.1 + %param_1.1556 = bf16[1,4096,8,128]{1,3,2,0:T(8,128)(2,1)S(1)} parameter(1) + %fusion.443 = bf16[4096,8,128]{0,2,1:T(8,128)(2,1)} fusion(%param_1.1556), kind=kLoop, calls=%fused_computation.105.clone.clone.clone.1, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/dot_general" stack_frame_id=0} + %convolution.153 = bf16[128,4096,1]{1,0,2:T(8,128)(2,1)} convolution(%fusion.444, %fusion.443), window={size=8}, dim_labels=b0f_o0i->bf0, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/dot_general" stack_frame_id=0} + ROOT %bitcast.855 = bf16[128,4096]{1,0:T(8,128)(2,1)S(1)} bitcast(%convolution.153), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/dot_general" stack_frame_id=0} +} + +%fused_computation.228.clone.clone.clone.1 (param_0.1504: bf16[8,128,128]) -> bf16[128,8,128] { + %param_0.1504 = bf16[8,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.860 = bf16[128,8,128]{2,0,1:T(8,128)(2,1)} bitcast(%param_0.1504) +} + +%fused_computation.101.clone.clone.clone.clone.clone.clone.clone.clone.clone.clone.clone.1 (param_0.1503: bf16[1,4096,8,128]) -> bf16[4096,8,128] { + %param_0.1503 = bf16[1,4096,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.859 = bf16[4096,8,128]{2,0,1:T(8,128)(2,1)} bitcast(%param_0.1503), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/dot_general" stack_frame_id=0} +} + +%region_16.20 (reduce_sum.129: f32[], reduce_sum.133: f32[]) -> f32[] { + %reduce_sum.129 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/reduce_sum"} + %reduce_sum.133 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/reduce_sum"} + ROOT %reduce_sum.134 = f32[]{:T(128)} add(%reduce_sum.129, %reduce_sum.133), metadata={op_name="checkpoint/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%fused_computation.155.clone.1 (param_0.1505: bf16[4,1,128,4096], param_1.1558: s32[], param_2.1190: bf16[4096], param_3.805: bf16[128,4096], param_4.495: bf16[128,4096], param_5.418: bf16[1,4096,8,128], param_6.269: bf16[8,128,128]) -> (f32[128], bf16[1,128,4096]) { + %param_0.1505 = bf16[4,1,128,4096]{3,2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1558 = s32[]{:T(128)S(6)} parameter(1) + %constant.362.clone.81 = s32[]{:T(128)} constant(0) + %dynamic-slice.185 = bf16[1,1,128,4096]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1505, %param_1.1558, %constant.362.clone.81, %constant.362.clone.81, %constant.362.clone.81), dynamic_slice_sizes={1,1,128,4096}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/jit(dynamic_index_in_dim)/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %bitcast.861 = bf16[1,128,4096]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic-slice.185), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/jit(dynamic_index_in_dim)/squeeze" stack_frame_id=0} + %convert_element_type.1040 = f32[1,128,4096]{2,1,0:T(8,128)} convert(%bitcast.861), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %param_6.269 = bf16[8,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(6) + %fusion.249.clone.3 = bf16[128,8,128]{2,0,1:T(8,128)(2,1)} fusion(%param_6.269), kind=kLoop, calls=%fused_computation.228.clone.clone.clone.1 + %param_5.418 = bf16[1,4096,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(5) + %fusion.167.clone.3 = bf16[4096,8,128]{2,0,1:T(8,128)(2,1)} fusion(%param_5.418), kind=kLoop, calls=%fused_computation.101.clone.clone.clone.clone.clone.clone.clone.clone.clone.clone.clone.1, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/dot_general" stack_frame_id=0} + %convolution.103.clone.3 = bf16[128,4096,1]{1,0,2:T(8,128)(2,1)} convolution(%fusion.249.clone.3, %fusion.167.clone.3), window={size=8}, dim_labels=b0f_o0i->bf0, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/dot_general" stack_frame_id=0} + %bitcast.337.clone.3 = bf16[128,4096]{1,0:T(8,128)(2,1)} bitcast(%convolution.103.clone.3), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/dot_general" stack_frame_id=0} + %convert.258 = f32[128,4096]{1,0:T(8,128)} convert(%bitcast.337.clone.3) + %param_4.495 = bf16[128,4096]{1,0:T(8,128)(2,1)S(1)} parameter(4) + %convert.259 = f32[128,4096]{1,0:T(8,128)} convert(%param_4.495) + %add_any.130.clone.3 = f32[128,4096]{1,0:T(8,128)} add(%convert.258, %convert.259), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/add_any" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["128","4096"],"tuple_shapes":[],"layout":{"minor_to_major":["1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false]}}} diff --git a/tests/utils/reference_hlo_qwen3_1.7b.txt b/tests/utils/reference_hlo_qwen3_1.7b.txt index f1ede66966..0337dd0166 100644 --- a/tests/utils/reference_hlo_qwen3_1.7b.txt +++ b/tests/utils/reference_hlo_qwen3_1.7b.txt @@ -1,4 +1,4 @@ -HloModule jit_train_step, is_scheduled=true, input_output_alias={ {0}: (0, {}, may-alias), {1}: (1, {}, may-alias), {2}: (2, {}, may-alias), {3}: (3, {}, may-alias), {4}: (4, {}, may-alias), {5}: (5, {}, may-alias), {6}: (6, {}, may-alias), {7}: (7, {}, may-alias), {8}: (8, {}, may-alias), {9}: (9, {}, may-alias), {10}: (10, {}, may-alias), {11}: (11, {}, may-alias), {12}: (12, {}, may-alias), {13}: (13, {}, may-alias), {14}: (14, {}, may-alias), {15}: (15, {}, may-alias), {16}: (16, {}, may-alias), {17}: (17, {}, may-alias), {18}: (18, {}, may-alias), {19}: (19, {}, may-alias), {20}: (20, {}, may-alias), {21}: (21, {}, may-alias), {22}: (22, {}, may-alias), {23}: (23, {}, may-alias), {24}: (24, {}, may-alias), {25}: (25, {}, may-alias), {26}: (26, {}, may-alias), {27}: (27, {}, may-alias), {28}: (28, {}, may-alias), {29}: (29, {}, may-alias), {30}: (30, {}, may-alias), {31}: (31, {}, may-alias), {32}: (32, {}, may-alias), {33}: (33, {}, may-alias), {34}: (34, {}, may-alias), {35}: (35, {}, may-alias), {36}: (36, {}, may-alias), {37}: (37, {}, may-alias), {38}: (38, {}, may-alias), {39}: (39, {}, may-alias), {40}: (40, {}, may-alias), {41}: (41, {}, may-alias) }, entry_computation_layout={(s32[]{:T(128)}, f32[2048]{0:T(1024)}, f32[2048,4,6144]{2,1,0:T(4,128)}, f32[2048,4,6144]{2,1,0:T(4,128)}, f32[6144,4,2048]{2,1,0:T(4,128)}, /*index=5*/f32[2048,4]{0,1:T(4,128)}, f32[2048,4]{0,1:T(4,128)}, f32[2048,4,8,128]{3,2,1,0:T(8,128)}, f32[128,4]{0,1:T(4,128)}, f32[16,4,128,2048]{3,2,1,0:T(8,128)}, /*index=10*/f32[2048,4,16,128]{3,2,1,0:T(8,128)}, f32[128,4]{0,1:T(4,128)}, f32[2048,4,8,128]{3,2,1,0:T(8,128)}, f32[151936,2048]{1,0:T(8,128)}, s32[]{:T(128)}, /*index=15*/f32[2048]{0:T(1024)}, f32[2048,4,6144]{2,1,0:T(4,128)}, f32[2048,4,6144]{2,1,0:T(4,128)}, f32[6144,4,2048]{2,1,0:T(4,128)}, f32[2048,4]{0,1:T(4,128)}, /*index=20*/f32[2048,4]{0,1:T(4,128)}, f32[2048,4,8,128]{3,2,1,0:T(8,128)}, f32[128,4]{0,1:T(4,128)}, f32[16,4,128,2048]{3,2,1,0:T(8,128)}, f32[2048,4,16,128]{3,2,1,0:T(8,128)}, /*index=25*/f32[128,4]{0,1:T(4,128)}, f32[2048,4,8,128]{3,2,1,0:T(8,128)}, f32[151936,2048]{1,0:T(8,128)}, f32[2048]{0:T(1024)}, f32[2048,4,6144]{2,1,0:T(4,128)}, /*index=30*/f32[2048,4,6144]{2,1,0:T(4,128)}, f32[6144,4,2048]{2,1,0:T(4,128)}, f32[2048,4]{0,1:T(4,128)}, f32[2048,4]{0,1:T(4,128)}, f32[2048,4,8,128]{3,2,1,0:T(8,128)}, /*index=35*/f32[128,4]{0,1:T(4,128)}, f32[16,4,128,2048]{3,2,1,0:T(8,128)}, f32[2048,4,16,128]{3,2,1,0:T(8,128)}, f32[128,4]{0,1:T(4,128)}, f32[2048,4,8,128]{3,2,1,0:T(8,128)}, /*index=40*/f32[151936,2048]{1,0:T(8,128)}, s32[]{:T(128)}, s32[4,128]{1,0:T(4,128)}, s32[4,128]{1,0:T(4,128)}, s32[4,128]{1,0:T(4,128)}, /*index=45*/s32[4,128]{1,0:T(4,128)}, s32[4,128]{1,0:T(4,128)})->(s32[]{:T(128)}, f32[2048]{0:T(1024)}, f32[2048,4,6144]{2,1,0:T(4,128)}, f32[2048,4,6144]{2,1,0:T(4,128)}, f32[6144,4,2048]{2,1,0:T(4,128)}, /*index=5*/f32[2048,4]{0,1:T(4,128)}, f32[2048,4]{0,1:T(4,128)}, f32[2048,4,8,128]{3,2,1,0:T(8,128)}, f32[128,4]{0,1:T(4,128)}, f32[16,4,128,2048]{3,2,1,0:T(8,128)}, /*index=10*/f32[2048,4,16,128]{3,2,1,0:T(8,128)}, f32[128,4]{0,1:T(4,128)}, f32[2048,4,8,128]{3,2,1,0:T(8,128)}, f32[151936,2048]{1,0:T(8,128)}, s32[]{:T(128)}, /*index=15*/f32[2048]{0:T(1024)}, f32[2048,4,6144]{2,1,0:T(4,128)}, f32[2048,4,6144]{2,1,0:T(4,128)}, f32[6144,4,2048]{2,1,0:T(4,128)}, f32[2048,4]{0,1:T(4,128)}, /*index=20*/f32[2048,4]{0,1:T(4,128)}, f32[2048,4,8,128]{3,2,1,0:T(8,128)}, f32[128,4]{0,1:T(4,128)}, f32[16,4,128,2048]{3,2,1,0:T(8,128)}, f32[2048,4,16,128]{3,2,1,0:T(8,128)}, /*index=25*/f32[128,4]{0,1:T(4,128)}, f32[2048,4,8,128]{3,2,1,0:T(8,128)}, f32[151936,2048]{1,0:T(8,128)}, f32[2048]{0:T(1024)}, f32[2048,4,6144]{2,1,0:T(4,128)}, /*index=30*/f32[2048,4,6144]{2,1,0:T(4,128)}, f32[6144,4,2048]{2,1,0:T(4,128)}, f32[2048,4]{0,1:T(4,128)}, f32[2048,4]{0,1:T(4,128)}, f32[2048,4,8,128]{3,2,1,0:T(8,128)}, /*index=35*/f32[128,4]{0,1:T(4,128)}, f32[16,4,128,2048]{3,2,1,0:T(8,128)}, f32[2048,4,16,128]{3,2,1,0:T(8,128)}, f32[128,4]{0,1:T(4,128)}, f32[2048,4,8,128]{3,2,1,0:T(8,128)}, /*index=40*/f32[151936,2048]{1,0:T(8,128)}, s32[]{:T(128)}, f32[]{:T(128)}, f32[]{:T(128)}, f32[]{:T(128)}, /*index=45*/f32[]{:T(128)}, f32[]{:T(128)}, f32[]{:T(128)}, f32[]{:T(128)}, f32[]{:T(128)}, /*index=50*/f32[]{:T(128)}, s32[]{:T(128)}, f32[]{:T(128)})}, allow_spmd_sharding_propagation_to_parameters={false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false}, allow_spmd_sharding_propagation_to_output={false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,true,true,true,true,true,true,true,true,true,true,true}, num_partitions=4 +HloModule jit_train_step, is_scheduled=true, input_output_alias={ {0}: (0, {}, may-alias), {1}: (1, {}, may-alias), {2}: (2, {}, may-alias), {3}: (3, {}, may-alias), {4}: (4, {}, may-alias), {5}: (5, {}, may-alias), {6}: (6, {}, may-alias), {7}: (7, {}, may-alias), {8}: (8, {}, may-alias), {9}: (9, {}, may-alias), {10}: (10, {}, may-alias), {11}: (11, {}, may-alias), {12}: (12, {}, may-alias), {13}: (13, {}, may-alias), {14}: (14, {}, may-alias), {15}: (15, {}, may-alias), {16}: (16, {}, may-alias), {17}: (17, {}, may-alias), {18}: (18, {}, may-alias), {19}: (19, {}, may-alias), {20}: (20, {}, may-alias), {21}: (21, {}, may-alias), {22}: (22, {}, may-alias), {23}: (23, {}, may-alias), {24}: (24, {}, may-alias), {25}: (25, {}, may-alias), {26}: (26, {}, may-alias), {27}: (27, {}, may-alias), {28}: (28, {}, may-alias), {29}: (29, {}, may-alias), {30}: (30, {}, may-alias), {31}: (31, {}, may-alias), {32}: (32, {}, may-alias), {33}: (33, {}, may-alias), {34}: (34, {}, may-alias), {35}: (35, {}, may-alias), {36}: (36, {}, may-alias), {37}: (37, {}, may-alias), {38}: (38, {}, may-alias), {39}: (39, {}, may-alias), {40}: (40, {}, may-alias), {41}: (41, {}, may-alias), {42}: (42, {}, may-alias), {43}: (43, {}, may-alias), {44}: (44, {}, may-alias), {45}: (45, {}, may-alias), {46}: (46, {}, may-alias), {47}: (47, {}, may-alias), {48}: (48, {}, may-alias), {49}: (49, {}, may-alias), {50}: (50, {}, may-alias), {51}: (51, {}, may-alias), {52}: (52, {}, may-alias), {53}: (53, {}, may-alias), {54}: (54, {}, may-alias), {55}: (55, {}, may-alias), {56}: (56, {}, may-alias), {57}: (57, {}, may-alias), {58}: (58, {}, may-alias), {59}: (59, {}, may-alias), {60}: (60, {}, may-alias), {61}: (61, {}, may-alias), {62}: (62, {}, may-alias), {63}: (63, {}, may-alias), {64}: (64, {}, may-alias), {65}: (65, {}, may-alias) }, entry_computation_layout={(f32[2048]{0:T(1024)}, u32[]{:T(128)}, u32[4]{0:T(128)}, u32[]{:T(128)}, u32[4]{0:T(128)}, /*index=5*/u32[]{:T(128)}, u32[4]{0:T(128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, /*index=10*/u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, f32[512,4,6144]{2,1,0:T(4,128)}, f32[512,4,6144]{2,1,0:T(4,128)}, /*index=15*/f32[6144,4,512]{2,1,0:T(4,128)}, f32[2048,4]{0,1:T(4,128)}, f32[2048,4]{0,1:T(4,128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, /*index=20*/u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, f32[512,4,8,128]{3,2,1,0:T(8,128)}, /*index=25*/f32[128,4]{0,1:T(4,128)}, f32[16,4,128,512]{3,2,1,0:T(8,128)}, f32[512,4,16,128]{3,2,1,0:T(8,128)}, f32[128,4]{0,1:T(4,128)}, f32[512,4,8,128]{3,2,1,0:T(8,128)}, /*index=30*/u32[]{:T(128)}, u32[4]{0:T(128)}, u32[]{:T(128)}, u32[4]{0:T(128)}, u32[]{:T(128)}, /*index=35*/u32[4]{0:T(128)}, f32[151936,512]{1,0:T(8,128)}, s32[]{:T(128)}, f32[2048]{0:T(1024)}, f32[512,4,6144]{2,1,0:T(4,128)}, /*index=40*/f32[512,4,6144]{2,1,0:T(4,128)}, f32[6144,4,512]{2,1,0:T(4,128)}, f32[2048,4]{0,1:T(4,128)}, f32[2048,4]{0,1:T(4,128)}, f32[512,4,8,128]{3,2,1,0:T(8,128)}, /*index=45*/f32[128,4]{0,1:T(4,128)}, f32[16,4,128,512]{3,2,1,0:T(8,128)}, f32[512,4,16,128]{3,2,1,0:T(8,128)}, f32[128,4]{0,1:T(4,128)}, f32[512,4,8,128]{3,2,1,0:T(8,128)}, /*index=50*/f32[151936,512]{1,0:T(8,128)}, f32[2048]{0:T(1024)}, f32[512,4,6144]{2,1,0:T(4,128)}, f32[512,4,6144]{2,1,0:T(4,128)}, f32[6144,4,512]{2,1,0:T(4,128)}, /*index=55*/f32[2048,4]{0,1:T(4,128)}, f32[2048,4]{0,1:T(4,128)}, f32[512,4,8,128]{3,2,1,0:T(8,128)}, f32[128,4]{0,1:T(4,128)}, f32[16,4,128,512]{3,2,1,0:T(8,128)}, /*index=60*/f32[512,4,16,128]{3,2,1,0:T(8,128)}, f32[128,4]{0,1:T(4,128)}, f32[512,4,8,128]{3,2,1,0:T(8,128)}, f32[151936,512]{1,0:T(8,128)}, s32[]{:T(128)}, /*index=65*/u32[]{:T(128)}, s32[1,128]{1,0:T(1,128)}, s32[1,128]{1,0:T(1,128)}, s32[1,128]{1,0:T(1,128)}, s32[1,128]{1,0:T(1,128)}, /*index=70*/s32[1,128]{1,0:T(1,128)})->(f32[2048]{0:T(1024)}, u32[]{:T(128)}, u32[4]{0:T(128)}, u32[]{:T(128)}, u32[4]{0:T(128)}, /*index=5*/u32[]{:T(128)}, u32[4]{0:T(128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, /*index=10*/u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, f32[512,4,6144]{2,1,0:T(4,128)}, f32[512,4,6144]{2,1,0:T(4,128)}, /*index=15*/f32[6144,4,512]{2,1,0:T(4,128)}, f32[2048,4]{0,1:T(4,128)}, f32[2048,4]{0,1:T(4,128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, /*index=20*/u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, f32[512,4,8,128]{3,2,1,0:T(8,128)}, /*index=25*/f32[128,4]{0,1:T(4,128)}, f32[16,4,128,512]{3,2,1,0:T(8,128)}, f32[512,4,16,128]{3,2,1,0:T(8,128)}, f32[128,4]{0,1:T(4,128)}, f32[512,4,8,128]{3,2,1,0:T(8,128)}, /*index=30*/u32[]{:T(128)}, u32[4]{0:T(128)}, u32[]{:T(128)}, u32[4]{0:T(128)}, u32[]{:T(128)}, /*index=35*/u32[4]{0:T(128)}, f32[151936,512]{1,0:T(8,128)}, s32[]{:T(128)}, f32[2048]{0:T(1024)}, f32[512,4,6144]{2,1,0:T(4,128)}, /*index=40*/f32[512,4,6144]{2,1,0:T(4,128)}, f32[6144,4,512]{2,1,0:T(4,128)}, f32[2048,4]{0,1:T(4,128)}, f32[2048,4]{0,1:T(4,128)}, f32[512,4,8,128]{3,2,1,0:T(8,128)}, /*index=45*/f32[128,4]{0,1:T(4,128)}, f32[16,4,128,512]{3,2,1,0:T(8,128)}, f32[512,4,16,128]{3,2,1,0:T(8,128)}, f32[128,4]{0,1:T(4,128)}, f32[512,4,8,128]{3,2,1,0:T(8,128)}, /*index=50*/f32[151936,512]{1,0:T(8,128)}, f32[2048]{0:T(1024)}, f32[512,4,6144]{2,1,0:T(4,128)}, f32[512,4,6144]{2,1,0:T(4,128)}, f32[6144,4,512]{2,1,0:T(4,128)}, /*index=55*/f32[2048,4]{0,1:T(4,128)}, f32[2048,4]{0,1:T(4,128)}, f32[512,4,8,128]{3,2,1,0:T(8,128)}, f32[128,4]{0,1:T(4,128)}, f32[16,4,128,512]{3,2,1,0:T(8,128)}, /*index=60*/f32[512,4,16,128]{3,2,1,0:T(8,128)}, f32[128,4]{0,1:T(4,128)}, f32[512,4,8,128]{3,2,1,0:T(8,128)}, f32[151936,512]{1,0:T(8,128)}, s32[]{:T(128)}, /*index=65*/u32[]{:T(128)}, f32[]{:T(128)}, f32[]{:T(128)}, f32[]{:T(128)}, f32[]{:T(128)}, /*index=70*/f32[]{:T(128)}, f32[]{:T(128)}, f32[]{:T(128)}, f32[]{:T(128)}, f32[]{:T(128)}, /*index=75*/s32[]{:T(128)}, f32[]{:T(128)})}, allow_spmd_sharding_propagation_to_parameters={false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false}, allow_spmd_sharding_propagation_to_output={false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,false,true,true,true,true,true,true,true,true,true,true,true}, num_partitions=4, frontend_attributes={arg_layout_modes="default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default",arg_memory_spaces="0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0",out_layout_modes="default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default;default",out_memory_spaces="0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0;0"} FileNames @@ -9,1992 +9,1992 @@ FileLocations StackFrames -%fused_computation (param_0.2: bf16[151936,2048], param_1.7: s32[1024]) -> bf16[512,2048] { - %param_0.2 = bf16[151936,2048]{1,0:T(8,128)(2,1)} parameter(0) - %param_1.7 = s32[1024]{0:T(1024)S(1)} parameter(1) - %custom-call.1 = s32[1024]{0:T(1024)} custom-call(%param_1.7), custom_call_target="AssumeGatherIndicesInBound", operand_layout_constraints={s32[1024]{0:T(1024)}}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} - %slice.6 = s32[512]{0:T(512)} slice(%custom-call.1), slice={[0:512]}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} - %reshape.444 = s32[4,128]{1,0:T(4,128)} reshape(%slice.6), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/broadcast_in_dim" stack_frame_id=0} - %transpose.461 = s32[4,128]{1,0:T(4,128)} transpose(%reshape.444), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/broadcast_in_dim" stack_frame_id=0} - %gather.4 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} gather(%param_0.2, %transpose.461), offset_dims={2}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=2, slice_sizes={1,2048}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} - %transpose.460 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} transpose(%gather.4), dimensions={0,1,2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} - ROOT %reshape.443 = bf16[512,2048]{1,0:T(8,128)(2,1)S(1)} reshape(%transpose.460), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} -} - -%region_42.47.clone (scatter-add.6: bf16[], scatter-add.7: bf16[]) -> bf16[] { - %scatter-add.7 = bf16[]{:T(256)} parameter(1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/scatter-add"} - %scatter-add.6 = bf16[]{:T(256)} parameter(0), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/scatter-add"} - ROOT %add.584 = bf16[]{:T(256)} add(%scatter-add.6, %scatter-add.7), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/add" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.1 (param_0.3: bf16[151936,2048], param_1.5: s32[512], param_2.4: bf16[512,2048]) -> bf16[151936,2048] { - %param_0.3 = bf16[151936,2048]{1,0:T(8,128)(2,1)} parameter(0) - %param_1.5 = s32[512]{0:T(512)S(1)} parameter(1) - %reshape.451 = s32[4,128]{1,0:T(4,128)} reshape(%param_1.5), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/broadcast_in_dim" stack_frame_id=0} - %transpose.466 = s32[4,128]{1,0:T(4,128)} transpose(%reshape.451), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/broadcast_in_dim" stack_frame_id=0} - %param_2.4 = bf16[512,2048]{1,0:T(8,128)(2,1)S(1)} parameter(2) - %reshape.452 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} reshape(%param_2.4), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while" stack_frame_id=0} - %transpose.467 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} transpose(%reshape.452), dimensions={0,1,2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while" stack_frame_id=0} - ROOT %scatter.2 = bf16[151936,2048]{1,0:T(8,128)(2,1)} scatter(%param_0.3, %transpose.466, %transpose.467), update_window_dims={2}, inserted_window_dims={0}, scatter_dims_to_operand_dims={0}, index_vector_dim=2, to_apply=%region_42.47.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/scatter-add" stack_frame_id=0} -} - -%region_71.76 (reduce_sum.464: f32[], reduce_sum.465: f32[]) -> f32[] { - %reduce_sum.465 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.464 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.466 = f32[]{:T(128)} add(%reduce_sum.464, %reduce_sum.465), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_56.61 (reduce_sum.386: f32[], reduce_sum.387: f32[]) -> f32[] { - %reduce_sum.387 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.386 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.388 = f32[]{:T(128)} add(%reduce_sum.386, %reduce_sum.387), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.277 (param_0.1368: f32[151936,2048], param_1.1556: f32[], param_2.1314: f32[], param_3.918: f32[], param_4.556: f32[151936,2048], param_5.468: f32[], param_6.358: bf16[151936,2048], param_7.201: bf16[151936,2048,1], param_8.118: pred[], param_9.97: f32[151936,2048]) -> (f32[], f32[151936,2048], f32[151936,2048], f32[151936,2048], f32[]) { - %param_0.1368 = f32[151936,2048]{1,0:T(8,128)} parameter(0) - %param_3.918 = f32[]{:T(128)S(6)} parameter(3) - %mul.1926.clone.1 = f32[151936,2048]{1,0:T(8,128)} broadcast(%param_3.918), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.118 = pred[]{:T(512)S(6)} parameter(8) - %select_n.268.clone.1 = pred[151936,2048]{1,0:T(8,128)(4,1)} broadcast(%param_8.118), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_7.201 = bf16[151936,2048,1]{1,0,2:T(8,128)(2,1)} parameter(7) - %bitcast.464.clone.1 = bf16[151936,2048]{1,0:T(8,128)(2,1)} bitcast(%param_7.201), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/dot_general" stack_frame_id=0} - %convert_element_type.1409.clone.1 = f32[151936,2048]{1,0:T(8,128)} convert(%bitcast.464.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %param_6.358 = bf16[151936,2048]{1,0:T(8,128)(2,1)} parameter(6) - %convert_element_type.1408.clone.1 = f32[151936,2048]{1,0:T(8,128)} convert(%param_6.358), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} - %add_any.197.clone.1 = f32[151936,2048]{1,0:T(8,128)} add(%convert_element_type.1409.clone.1, %convert_element_type.1408.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/add_any" stack_frame_id=0} - %param_5.468 = f32[]{:T(128)} parameter(5) - %div.860.clone.1 = f32[151936,2048]{1,0:T(8,128)} broadcast(%param_5.468), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.859.clone.1 = f32[151936,2048]{1,0:T(8,128)} divide(%add_any.197.clone.1, %div.860.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.267.clone.1 = f32[151936,2048]{1,0:T(8,128)} select(%select_n.268.clone.1, %add_any.197.clone.1, %div.859.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %constant.1092.clone.1 = f32[]{:T(128)} constant(0.1) - %broadcast.844.clone.1 = f32[151936,2048]{1,0:T(8,128)} broadcast(%constant.1092.clone.1), dimensions={}, metadata={op_name="broadcast.74"} - %mul.1932.clone.1 = f32[151936,2048]{1,0:T(8,128)} multiply(%select_n.267.clone.1, %broadcast.844.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_9.97 = f32[151936,2048]{1,0:T(8,128)} parameter(9) - %constant.1096.clone.1 = f32[]{:T(128)} constant(0.9) - %mul.1933.clone.1 = f32[151936,2048]{1,0:T(8,128)} broadcast(%constant.1096.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1931.clone.1 = f32[151936,2048]{1,0:T(8,128)} multiply(%param_9.97, %mul.1933.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.941.clone.1 = f32[151936,2048]{1,0:T(8,128)} add(%mul.1932.clone.1, %mul.1931.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1314 = f32[]{:T(128)S(6)} parameter(2) - %div.856.clone.1 = f32[151936,2048]{1,0:T(8,128)} broadcast(%param_2.1314), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %integer_pow.65.clone.1 = f32[151936,2048]{1,0:T(8,128)} multiply(%select_n.267.clone.1, %select_n.267.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} - %constant.1095.clone.1 = f32[]{:T(128)} constant(0.05) - %mul.1930.clone.1 = f32[151936,2048]{1,0:T(8,128)} broadcast(%constant.1095.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1928.clone.1 = f32[151936,2048]{1,0:T(8,128)} multiply(%integer_pow.65.clone.1, %mul.1930.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.556 = f32[151936,2048]{1,0:T(8,128)} parameter(4) - %constant.1094.clone.1 = f32[]{:T(128)} constant(0.95) - %mul.1929.clone.1 = f32[151936,2048]{1,0:T(8,128)} broadcast(%constant.1094.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1927.clone.1 = f32[151936,2048]{1,0:T(8,128)} multiply(%param_4.556, %mul.1929.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.940.clone.1 = f32[151936,2048]{1,0:T(8,128)} add(%mul.1928.clone.1, %mul.1927.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1556 = f32[]{:T(128)S(6)} parameter(1) - %div.855.clone.1 = f32[151936,2048]{1,0:T(8,128)} broadcast(%param_1.1556), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.854.clone.1 = f32[151936,2048]{1,0:T(8,128)} divide(%add.940.clone.1, %div.855.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %sqrt.62.clone.1 = f32[151936,2048]{1,0:T(8,128)} sqrt(%div.854.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} - %constant.1093.clone.1 = f32[]{:T(128)} constant(1e-08) - %add.939.clone.1 = f32[151936,2048]{1,0:T(8,128)} broadcast(%constant.1093.clone.1), dimensions={}, metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %add.938.clone.1 = f32[151936,2048]{1,0:T(8,128)} add(%sqrt.62.clone.1, %add.939.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %multiply.426.clone.1 = f32[151936,2048]{1,0:T(8,128)} multiply(%div.856.clone.1, %add.938.clone.1), metadata={op_name="multiply.61"} - %div.853.clone.1 = f32[151936,2048]{1,0:T(8,128)} divide(%add.941.clone.1, %multiply.426.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1925.clone.1 = f32[151936,2048]{1,0:T(8,128)} multiply(%param_0.1368, %broadcast.844.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.937.clone.1 = f32[151936,2048]{1,0:T(8,128)} add(%div.853.clone.1, %mul.1925.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1924.clone.1 = f32[151936,2048]{1,0:T(8,128)} multiply(%mul.1926.clone.1, %add.937.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.936.clone.1 = f32[151936,2048]{1,0:T(8,128)} add(%param_0.1368, %mul.1924.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.214 = f32[151936,2048]{1,0:T(8,128)} multiply(%add.936.clone.1, %add.936.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1200 = f32[]{:T(128)} constant(0) - %reduce.176 = f32[]{:T(128)} reduce(%square.214, %constant.1200), dimensions={0,1}, to_apply=%region_71.76, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.178.clone.1 = f32[]{:T(128)} reduce(%integer_pow.65.clone.1, %constant.1200), dimensions={0,1}, to_apply=%region_56.61, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.144 = (f32[]{:T(128)}, f32[151936,2048]{1,0:T(8,128)}, f32[151936,2048]{1,0:T(8,128)}, f32[151936,2048]{1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.176, %add.936.clone.1, %add.940.clone.1, %add.941.clone.1, %reduce.178.clone.1) -} - -%region_43.48 (reduce_sum.317: f32[], reduce_sum.318: f32[]) -> f32[] { - %reduce_sum.318 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.317 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.319 = f32[]{:T(128)} add(%reduce_sum.317, %reduce_sum.318), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.367.clone.clone (param_0.1355: f32[4,128], param_1.1549: bf16[4,128,2048], param_2.1290: bf16[2048]) -> bf16[4,128,2048] { - %param_2.1290 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(2) - %dot_general.480 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1290), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1549 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %convert_element_type.1451 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%param_1.1549), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %param_0.1355 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.2083 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_0.1355), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %mul.2082 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1451, %mul.2083), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %convert_element_type.1450 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.2082), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - ROOT %dot_general.479 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.480, %convert_element_type.1450), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} -} - -%fused_computation.289.clone.clone.clone (param_0.1356: bf16[4,128,151936], param_1.1550: s32[4,128], param_2.1291: f32[4,128], param_3.911: f32[4,128], param_4.546: bf16[4,128], param_5.446: f32[4,128]) -> bf16[4,128,151936] { - %param_5.446 = f32[4,128]{1,0:T(4,128)S(1)} parameter(5) - %mul.2087 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_5.446), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - %param_3.911 = f32[4,128]{1,0:T(4,128)S(1)} parameter(3) - %mul.2086 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_3.911), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - %param_0.1356 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} parameter(0) - %convert_element_type.1454 = f32[4,128,151936]{2,1,0:T(8,128)} convert(%param_0.1356), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %param_4.546 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(4) - %sub.94 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_4.546), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %sub.93 = f32[4,128,151936]{2,1,0:T(8,128)} subtract(%convert_element_type.1454, %sub.94), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %exp.62 = f32[4,128,151936]{2,1,0:T(8,128)} exponential(%sub.93), metadata={op_name="jit(train_step)/jvp()/exp" stack_frame_id=0} - %mul.2085 = f32[4,128,151936]{2,1,0:T(8,128)} multiply(%mul.2086, %exp.62), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - %param_2.1291 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %div.966 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_2.1291), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} - %div.965 = f32[4,128,151936]{2,1,0:T(8,128)} divide(%mul.2085, %div.966), metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} - %param_1.1550 = s32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %eq.49 = s32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_1.1550), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} - %eq.48 = s32[4,128,151936]{2,1,0:T(8,128)} iota(), iota_dimension=2, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} - %eq.47 = pred[4,128,151936]{2,1,0:T(8,128)(4,1)} compare(%eq.49, %eq.48), direction=EQ, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} - %convert_element_type.1453 = f32[4,128,151936]{2,1,0:T(8,128)} convert(%eq.47), metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/convert_element_type" stack_frame_id=0} - %sub.92 = f32[4,128,151936]{2,1,0:T(8,128)} subtract(%div.965, %convert_element_type.1453), metadata={op_name="jit(train_step)/transpose(jvp())/sub" stack_frame_id=0} - %mul.2084 = f32[4,128,151936]{2,1,0:T(8,128)} multiply(%mul.2087, %sub.92), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - ROOT %convert_element_type.1452 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} convert(%mul.2084), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} -} - -%fused_computation.281 (param_0.1381: bf16[151936,2048], param_1.1569: f32[4,128], param_2.1327: bf16[4,128,2048], param_3.931: bf16[2048], param_4.569: bf16[4,128,151936], param_5.481: s32[4,128], param_6.371: f32[4,128], param_7.214: f32[4,128], param_8.131: bf16[4,128], param_9.98: f32[4,128]) -> (f32[], bf16[151936,2048,1]) { - %param_4.569 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} parameter(4) - %param_5.481 = s32[4,128]{1,0:T(4,128)S(1)} parameter(5) - %param_6.371 = f32[4,128]{1,0:T(4,128)S(1)} parameter(6) - %param_7.214 = f32[4,128]{1,0:T(4,128)S(1)} parameter(7) - %param_8.131 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(8) - %param_9.98 = f32[4,128]{1,0:T(4,128)S(1)} parameter(9) - %multiply_convert_fusion.1.clone.1 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} fusion(%param_4.569, %param_5.481, %param_6.371, %param_7.214, %param_8.131, /*index=5*/%param_9.98), kind=kLoop, calls=%fused_computation.289.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %param_1.1569 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %param_2.1327 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) - %param_3.931 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(3) - %fusion.269.clone.1 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} fusion(%param_1.1569, %param_2.1327, %param_3.931), kind=kLoop, calls=%fused_computation.367.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %convolution.86.clone.1 = bf16[151936,2048,1]{1,0,2:T(8,128)(2,1)} convolution(%multiply_convert_fusion.1.clone.1, %fusion.269.clone.1), window={size=4}, dim_labels=0fb_0io->bf0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/dot_general" stack_frame_id=0} - %bitcast.333 = bf16[151936,2048]{1,0:T(8,128)(2,1)} bitcast(%convolution.86.clone.1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/dot_general" stack_frame_id=0} - %convert_element_type.1323 = f32[151936,2048]{1,0:T(8,128)} convert(%bitcast.333), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %param_0.1381 = bf16[151936,2048]{1,0:T(8,128)(2,1)} parameter(0) - %convert_element_type.1322 = f32[151936,2048]{1,0:T(8,128)} convert(%param_0.1381), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} - %add_any.184 = f32[151936,2048]{1,0:T(8,128)} add(%convert_element_type.1323, %convert_element_type.1322), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/add_any" stack_frame_id=0} - %square.215 = f32[151936,2048]{1,0:T(8,128)} multiply(%add_any.184, %add_any.184), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1213 = f32[]{:T(128)} constant(0) - %reduce.177 = f32[]{:T(128)} reduce(%square.215, %constant.1213), dimensions={0,1}, to_apply=%region_43.48, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.166 = (f32[]{:T(128)}, bf16[151936,2048,1]{1,0,2:T(8,128)(2,1)}) tuple(%reduce.177, %convolution.86.clone.1) -} - -%region_57.62 (reduce_sum.389: f32[], reduce_sum.393: f32[]) -> f32[] { - %reduce_sum.393 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} - %reduce_sum.389 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} - ROOT %reduce_sum.394 = f32[]{:T(128)} add(%reduce_sum.389, %reduce_sum.393), metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +%fused_computation.478 (param_0.1336: s32[1,128]) -> s32[1,1,128] { + %param_0.1336 = s32[1,128]{1,0:T(1,128)S(1)} parameter(0) + %constant.449.clone.9 = s32[]{:T(128)} constant(0) + %broadcast.982 = s32[1,128]{1,0:T(1,128)} broadcast(%constant.449.clone.9), dimensions={}, metadata={op_name="broadcast.99"} + %lt.32 = pred[1,128]{1,0:T(4,128)(4,1)} compare(%param_0.1336, %broadcast.982), direction=LT, metadata={op_name="jit(train_step)/jvp()/lt" stack_frame_id=0} + %constant.461.clone.1 = s32[]{:T(128)} constant(151936) + %add.941 = s32[1,128]{1,0:T(1,128)} broadcast(%constant.461.clone.1), dimensions={}, metadata={op_name="jit(train_step)/jvp()/add" stack_frame_id=0} + %add.930 = s32[1,128]{1,0:T(1,128)} add(%param_0.1336, %add.941), metadata={op_name="jit(train_step)/jvp()/add" stack_frame_id=0} + %select_n.180 = s32[1,128]{1,0:T(1,128)} select(%lt.32, %add.930, %param_0.1336), metadata={op_name="jit(train_step)/jvp()/select_n" stack_frame_id=0} + ROOT %bitcast.616 = s32[1,1,128]{2,1,0:T(1,128)S(1)} bitcast(%select_n.180) } -%fused_computation.288 (param_0.1392: bf16[4,128,151936], param_1.1577: f32[4,128], param_2.1330: s32[4,128], param_3.933: bf16[4,128]) -> f32[4,128] { - %param_2.1330 = s32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %eq.30 = s32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_2.1330), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} - %eq.25 = s32[4,128,151936]{2,1,0:T(8,128)} iota(), iota_dimension=2, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} - %eq.24 = pred[4,128,151936]{2,1,0:T(8,128)(4,1)} compare(%eq.30, %eq.25), direction=EQ, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} - %param_0.1392 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} parameter(0) - %convert_element_type.1340 = f32[4,128,151936]{2,1,0:T(8,128)} convert(%param_0.1392), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %param_3.933 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(3) - %sub.73 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_3.933), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %sub.64 = f32[4,128,151936]{2,1,0:T(8,128)} subtract(%convert_element_type.1340, %sub.73), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %param_1.1577 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %sub.71 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_1.1577), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %sub.60 = f32[4,128,151936]{2,1,0:T(8,128)} subtract(%sub.64, %sub.71), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %constant.1225 = f32[]{:T(128)} constant(0) - %broadcast.769 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%constant.1225), dimensions={}, metadata={op_name="broadcast.109"} - %mul.1765 = f32[4,128,151936]{2,1,0:T(8,128)} select(%eq.24, %sub.60, %broadcast.769), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} - ROOT %reduce.179 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%mul.1765, %constant.1225), dimensions={2}, to_apply=%region_57.62, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} +%fused_computation.448 (param_0.1269: s32[512]) -> s32[1024] { + %constant.1055 = s32[] constant(0), metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} + %broadcast.952 = s32[1024]{0:T(1024)} broadcast(%constant.1055), dimensions={}, metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} + %param_0.1269 = s32[512]{0:T(512)S(1)} parameter(0) + %constant.1056 = s32[] constant(2147483647), metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} + %pad.57 = s32[1024]{0:T(1024)} pad(%param_0.1269, %constant.1056), padding=0_512, metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} + %constant.1054 = s32[] constant(151935), metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} + %broadcast.951 = s32[1024]{0:T(1024)} broadcast(%constant.1054), dimensions={}, metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} + ROOT %clamp.3 = s32[1024]{0:T(1024)S(1)} clamp(%broadcast.952, %pad.57, %broadcast.951), metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} } -%region_9.12 (reduce_sum.186: f32[], reduce_sum.190: f32[]) -> f32[] { - %reduce_sum.190 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} - %reduce_sum.186 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} - ROOT %reduce_sum.191 = f32[]{:T(128)} add(%reduce_sum.186, %reduce_sum.190), metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +%fused_computation (param_0.2: bf16[151936,512], param_1.7: s32[1024]) -> bf16[512,512] { + %param_0.2 = bf16[151936,512]{1,0:T(8,128)(2,1)} parameter(0) + %param_1.7 = s32[1024]{0:T(1024)S(1)} parameter(1) + %custom-call.1 = s32[1024]{0:T(1024)} custom-call(%param_1.7), custom_call_target="AssumeGatherIndicesInBound", operand_layout_constraints={s32[1024]{0:T(1024)}}, metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} + %slice.34 = s32[512]{0:T(512)} slice(%custom-call.1), slice={[0:512]}, metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} + %reshape.793 = s32[4,128]{1,0:T(4,128)} reshape(%slice.34), metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} + %transpose.442 = s32[4,128]{1,0:T(4,128)} transpose(%reshape.793), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} + %gather.4 = bf16[4,128,512]{2,1,0:T(8,128)(2,1)} gather(%param_0.2, %transpose.442), offset_dims={2}, collapsed_slice_dims={0}, start_index_map={0}, index_vector_dim=2, slice_sizes={1,512}, metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} + %transpose.441 = bf16[4,128,512]{2,1,0:T(8,128)(2,1)} transpose(%gather.4), dimensions={0,1,2}, metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} + ROOT %reshape.792 = bf16[512,512]{1,0:T(8,128)(2,1)S(1)} reshape(%transpose.441), metadata={op_name="jit(train_step)/jvp()/gather" stack_frame_id=0} +} + +%fused_computation.494 (param_0.1517: f32[128,4]) -> bf16[4,128] { + %param_0.1517 = f32[128,4]{0,1:T(4,128)S(1)} parameter(0) + %bitcast.757 = f32[4,128]{1,0:T(4,128)} bitcast(%param_0.1517), metadata={op_name="jit(train_step)/jvp()/transpose" stack_frame_id=0} + ROOT %convert.178 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} convert(%bitcast.757) } -%fused_computation.293 (param_0.1393: bf16[4,128,151936], param_1.1578: bf16[4,128]) -> f32[4,128] { - %param_0.1393 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} parameter(0) - %convert_element_type.1346 = f32[4,128,151936]{2,1,0:T(8,128)} convert(%param_0.1393), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %param_1.1578 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(1) - %sub.74 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_1.1578), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %sub.70 = f32[4,128,151936]{2,1,0:T(8,128)} subtract(%convert_element_type.1346, %sub.74), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %exp.54 = f32[4,128,151936]{2,1,0:T(8,128)} exponential(%sub.70), metadata={op_name="jit(train_step)/jvp()/exp" stack_frame_id=0} - %constant.1226 = f32[]{:T(128)} constant(0) - ROOT %reduce.180 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%exp.54, %constant.1226), dimensions={2}, to_apply=%region_9.12, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} -} - -%region_33.38 (reduce_sum.269: f32[], reduce_sum.270: f32[]) -> f32[] { - %reduce_sum.270 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.269 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.274 = f32[]{:T(128)} add(%reduce_sum.269, %reduce_sum.270), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.298 (param_0.1387: f32[4,6144,2048]) -> f32[] { - %param_0.1387 = f32[4,6144,2048]{2,0,1:T(4,128)} parameter(0) - %bitcast.347 = f32[6144,4,2048]{2,1,0:T(4,128)} bitcast(%param_0.1387), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %square.218 = f32[6144,4,2048]{2,1,0:T(4,128)} multiply(%bitcast.347, %bitcast.347), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1219 = f32[]{:T(128)} constant(0) - ROOT %reduce.181 = f32[]{:T(128)} reduce(%square.218, %constant.1219), dimensions={0,1,2}, to_apply=%region_33.38, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} -} - -%region_32.37 (reduce_sum.263: f32[], reduce_sum.267: f32[]) -> f32[] { - %reduce_sum.267 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.263 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.268 = f32[]{:T(128)} add(%reduce_sum.263, %reduce_sum.267), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_31.36 (reduce_sum.260: f32[], reduce_sum.261: f32[]) -> f32[] { - %reduce_sum.261 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.260 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.262 = f32[]{:T(128)} add(%reduce_sum.260, %reduce_sum.261), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.300 (param_0.1388: f32[4,2048,6144], param_1.1573: f32[4,2048,6144]) -> (f32[], f32[]) { - %param_0.1388 = f32[4,2048,6144]{2,0,1:T(4,128)} parameter(0) - %bitcast.351 = f32[2048,4,6144]{2,1,0:T(4,128)} bitcast(%param_0.1388), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %square.221 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%bitcast.351, %bitcast.351), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1220 = f32[]{:T(128)} constant(0) - %reduce.182 = f32[]{:T(128)} reduce(%square.221, %constant.1220), dimensions={0,1,2}, to_apply=%region_32.37, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %param_1.1573 = f32[4,2048,6144]{2,0,1:T(4,128)} parameter(1) - %bitcast.355.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} bitcast(%param_1.1573), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %square.224.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%bitcast.355.clone.1, %bitcast.355.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %reduce.183.clone.1 = f32[]{:T(128)} reduce(%square.224.clone.1, %constant.1220), dimensions={0,1,2}, to_apply=%region_31.36, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.167 = (f32[]{:T(128)}, f32[]{:T(128)}) tuple(%reduce.182, %reduce.183.clone.1) -} - -%fused_computation.303 (param_0.901: f32[6144,4,2048]) -> bf16[4,6144,2048] { - %param_0.901 = f32[6144,4,2048]{2,1,0:T(4,128)} parameter(0) - %copy.190 = bf16[6144,4,2048]{2,0,1:T(8,128)(2,1)} copy(%param_0.901), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'mlp\'][\'wo\'][\'kernel\']"} - ROOT %bitcast.356 = bf16[4,6144,2048]{2,1,0:T(8,128)(2,1)} bitcast(%copy.190), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} -} - -%fused_computation.304 (param_0.903: f32[2048,4,6144]) -> bf16[4,2048,6144] { - %param_0.903 = f32[2048,4,6144]{2,1,0:T(4,128)} parameter(0) - %copy.191 = bf16[2048,4,6144]{2,0,1:T(8,128)(2,1)} copy(%param_0.903), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'mlp\'][\'wi_1\'][\'kernel\']"} - ROOT %bitcast.357 = bf16[4,2048,6144]{2,1,0:T(8,128)(2,1)} bitcast(%copy.191), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} -} - -%fused_computation.305 (param_0.905: f32[2048,4,6144]) -> bf16[4,2048,6144] { - %param_0.905 = f32[2048,4,6144]{2,1,0:T(4,128)} parameter(0) - %copy.192 = bf16[2048,4,6144]{2,0,1:T(8,128)(2,1)} copy(%param_0.905), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'mlp\'][\'wi_0\'][\'kernel\']"} - ROOT %bitcast.358 = bf16[4,2048,6144]{2,1,0:T(8,128)(2,1)} bitcast(%copy.192), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} -} - -%region_62.67 (reduce_sum.416: f32[], reduce_sum.417: f32[]) -> f32[] { - %reduce_sum.417 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.416 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.421 = f32[]{:T(128)} add(%reduce_sum.416, %reduce_sum.417), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_47.52 (reduce_sum.338: f32[], reduce_sum.339: f32[]) -> f32[] { - %reduce_sum.339 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.338 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.340 = f32[]{:T(128)} add(%reduce_sum.338, %reduce_sum.339), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.306 (param_0.1377: f32[6144,4,2048], param_1.1565: f32[], param_2.1323: f32[], param_3.927: f32[], param_4.565: f32[6144,4,2048], param_5.477: f32[], param_6.367: f32[4,6144,2048], param_7.210: pred[], param_8.127: f32[6144,4,2048]) -> (f32[], f32[6144,4,2048], f32[6144,4,2048], f32[6144,4,2048], f32[]) { - %param_0.1377 = f32[6144,4,2048]{2,1,0:T(4,128)} parameter(0) - %param_3.927 = f32[]{:T(128)S(6)} parameter(3) - %mul.1998.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} broadcast(%param_3.927), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.210 = pred[]{:T(512)S(6)} parameter(7) - %select_n.304.clone.1 = pred[6144,4,2048]{2,1,0:T(4,128)(4,1)} broadcast(%param_7.210), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.367 = f32[4,6144,2048]{2,0,1:T(4,128)} parameter(6) - %bitcast.482.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} bitcast(%param_6.367), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.477 = f32[]{:T(128)} parameter(5) - %div.932.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} broadcast(%param_5.477), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.931.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} divide(%bitcast.482.clone.1, %div.932.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.303.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} select(%select_n.304.clone.1, %bitcast.482.clone.1, %div.931.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %constant.1146.clone.1 = f32[]{:T(128)} constant(0.1) - %broadcast.886.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} broadcast(%constant.1146.clone.1), dimensions={}, metadata={op_name="broadcast.83"} - %mul.2004.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} multiply(%select_n.303.clone.1, %broadcast.886.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.127 = f32[6144,4,2048]{2,1,0:T(4,128)} parameter(8) - %constant.1150.clone.1 = f32[]{:T(128)} constant(0.9) - %mul.2005.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} broadcast(%constant.1150.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.2003.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} multiply(%param_8.127, %mul.2005.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.989.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} add(%mul.2004.clone.1, %mul.2003.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1323 = f32[]{:T(128)S(6)} parameter(2) - %div.928.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} broadcast(%param_2.1323), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %integer_pow.74.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} multiply(%select_n.303.clone.1, %select_n.303.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} - %constant.1149.clone.1 = f32[]{:T(128)} constant(0.05) - %mul.2002.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} broadcast(%constant.1149.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.2000.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} multiply(%integer_pow.74.clone.1, %mul.2002.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.565 = f32[6144,4,2048]{2,1,0:T(4,128)} parameter(4) - %constant.1148.clone.1 = f32[]{:T(128)} constant(0.95) - %mul.2001.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} broadcast(%constant.1148.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1999.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} multiply(%param_4.565, %mul.2001.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.988.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} add(%mul.2000.clone.1, %mul.1999.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1565 = f32[]{:T(128)S(6)} parameter(1) - %div.927.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} broadcast(%param_1.1565), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.926.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} divide(%add.988.clone.1, %div.927.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %sqrt.71.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} sqrt(%div.926.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} - %constant.1147.clone.1 = f32[]{:T(128)} constant(1e-08) - %add.987.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} broadcast(%constant.1147.clone.1), dimensions={}, metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %add.986.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} add(%sqrt.71.clone.1, %add.987.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %multiply.435.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} multiply(%div.928.clone.1, %add.986.clone.1), metadata={op_name="multiply.52"} - %div.925.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} divide(%add.989.clone.1, %multiply.435.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1997.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} multiply(%param_0.1377, %broadcast.886.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.985.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} add(%div.925.clone.1, %mul.1997.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1996.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} multiply(%mul.1998.clone.1, %add.985.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.984.clone.1 = f32[6144,4,2048]{2,1,0:T(4,128)} add(%param_0.1377, %mul.1996.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.225 = f32[6144,4,2048]{2,1,0:T(4,128)} multiply(%add.984.clone.1, %add.984.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1209 = f32[]{:T(128)} constant(0) - %reduce.184 = f32[]{:T(128)} reduce(%square.225, %constant.1209), dimensions={0,1,2}, to_apply=%region_62.67, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.187.clone.1 = f32[]{:T(128)} reduce(%integer_pow.74.clone.1, %constant.1209), dimensions={0,1,2}, to_apply=%region_47.52, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.145 = (f32[]{:T(128)}, f32[6144,4,2048]{2,1,0:T(4,128)}, f32[6144,4,2048]{2,1,0:T(4,128)}, f32[6144,4,2048]{2,1,0:T(4,128)}, f32[]{:T(128)}) tuple(%reduce.184, %add.984.clone.1, %add.988.clone.1, %add.989.clone.1, %reduce.187.clone.1) -} - -%region_61.66 (reduce_sum.410: f32[], reduce_sum.414: f32[]) -> f32[] { - %reduce_sum.414 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.410 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.415 = f32[]{:T(128)} add(%reduce_sum.410, %reduce_sum.414), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_46.51 (reduce_sum.332: f32[], reduce_sum.333: f32[]) -> f32[] { - %reduce_sum.333 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.332 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.337 = f32[]{:T(128)} add(%reduce_sum.332, %reduce_sum.333), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.307 (param_0.1378: f32[2048,4,6144], param_1.1566: f32[], param_2.1324: f32[], param_3.928: f32[], param_4.566: f32[2048,4,6144], param_5.478: f32[], param_6.368: f32[4,2048,6144], param_7.211: pred[], param_8.128: f32[2048,4,6144]) -> (f32[], f32[2048,4,6144], f32[2048,4,6144], f32[2048,4,6144], f32[]) { - %param_0.1378 = f32[2048,4,6144]{2,1,0:T(4,128)} parameter(0) - %param_3.928 = f32[]{:T(128)S(6)} parameter(3) - %mul.2008.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%param_3.928), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.211 = pred[]{:T(512)S(6)} parameter(7) - %select_n.308.clone.1 = pred[2048,4,6144]{2,1,0:T(4,128)(4,1)} broadcast(%param_7.211), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.368 = f32[4,2048,6144]{2,0,1:T(4,128)} parameter(6) - %bitcast.484.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} bitcast(%param_6.368), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.478 = f32[]{:T(128)} parameter(5) - %div.940.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%param_5.478), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.939.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} divide(%bitcast.484.clone.1, %div.940.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.307.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} select(%select_n.308.clone.1, %bitcast.484.clone.1, %div.939.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %constant.1152.clone.1 = f32[]{:T(128)} constant(0.1) - %broadcast.892.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%constant.1152.clone.1), dimensions={}, metadata={op_name="broadcast.85"} - %mul.2012.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%select_n.307.clone.1, %broadcast.892.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.128 = f32[2048,4,6144]{2,1,0:T(4,128)} parameter(8) - %constant.1156.clone.1 = f32[]{:T(128)} constant(0.9) - %broadcast.891.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%constant.1156.clone.1), dimensions={}, metadata={op_name="broadcast.84"} - %mul.2011.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%param_8.128, %broadcast.891.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.994.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} add(%mul.2012.clone.1, %mul.2011.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1324 = f32[]{:T(128)S(6)} parameter(2) - %div.936.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%param_2.1324), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %integer_pow.75.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%select_n.307.clone.1, %select_n.307.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} - %constant.1155.clone.1 = f32[]{:T(128)} constant(0.05) - %broadcast.890.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%constant.1155.clone.1), dimensions={}, metadata={op_name="broadcast.73"} - %mul.2010.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%integer_pow.75.clone.1, %broadcast.890.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.566 = f32[2048,4,6144]{2,1,0:T(4,128)} parameter(4) - %constant.1154.clone.1 = f32[]{:T(128)} constant(0.95) - %broadcast.889.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%constant.1154.clone.1), dimensions={}, metadata={op_name="broadcast.72"} - %mul.2009.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%param_4.566, %broadcast.889.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.993.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} add(%mul.2010.clone.1, %mul.2009.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1566 = f32[]{:T(128)S(6)} parameter(1) - %div.935.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%param_1.1566), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.934.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} divide(%add.993.clone.1, %div.935.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %sqrt.72.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} sqrt(%div.934.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} - %constant.1153.clone.1 = f32[]{:T(128)} constant(1e-08) - %broadcast.887.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%constant.1153.clone.1), dimensions={}, metadata={op_name="broadcast.65"} - %add.992.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} add(%sqrt.72.clone.1, %broadcast.887.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %multiply.436.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%div.936.clone.1, %add.992.clone.1), metadata={op_name="multiply.51"} - %div.933.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} divide(%add.994.clone.1, %multiply.436.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.2007.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%param_0.1378, %broadcast.892.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.991.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} add(%div.933.clone.1, %mul.2007.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.2006.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%mul.2008.clone.1, %add.991.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.990.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} add(%param_0.1378, %mul.2006.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.226 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%add.990.clone.1, %add.990.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1210 = f32[]{:T(128)} constant(0) - %reduce.185 = f32[]{:T(128)} reduce(%square.226, %constant.1210), dimensions={0,1,2}, to_apply=%region_61.66, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.188.clone.1 = f32[]{:T(128)} reduce(%integer_pow.75.clone.1, %constant.1210), dimensions={0,1,2}, to_apply=%region_46.51, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.146 = (f32[]{:T(128)}, f32[2048,4,6144]{2,1,0:T(4,128)}, f32[2048,4,6144]{2,1,0:T(4,128)}, f32[2048,4,6144]{2,1,0:T(4,128)}, f32[]{:T(128)}) tuple(%reduce.185, %add.990.clone.1, %add.993.clone.1, %add.994.clone.1, %reduce.188.clone.1) -} - -%region_60.65 (reduce_sum.407: f32[], reduce_sum.408: f32[]) -> f32[] { - %reduce_sum.408 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.407 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.409 = f32[]{:T(128)} add(%reduce_sum.407, %reduce_sum.408), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_45.50 (reduce_sum.326: f32[], reduce_sum.330: f32[]) -> f32[] { - %reduce_sum.330 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.326 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.331 = f32[]{:T(128)} add(%reduce_sum.326, %reduce_sum.330), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.308 (param_0.1379: f32[2048,4,6144], param_1.1567: f32[], param_2.1325: f32[], param_3.929: f32[], param_4.567: f32[2048,4,6144], param_5.479: f32[], param_6.369: f32[4,2048,6144], param_7.212: pred[], param_8.129: f32[2048,4,6144]) -> (f32[], f32[2048,4,6144], f32[2048,4,6144], f32[2048,4,6144], f32[]) { - %param_0.1379 = f32[2048,4,6144]{2,1,0:T(4,128)} parameter(0) - %param_3.929 = f32[]{:T(128)S(6)} parameter(3) - %mul.2015.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%param_3.929), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.212 = pred[]{:T(512)S(6)} parameter(7) - %select_n.312.clone.1 = pred[2048,4,6144]{2,1,0:T(4,128)(4,1)} broadcast(%param_7.212), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.369 = f32[4,2048,6144]{2,0,1:T(4,128)} parameter(6) - %bitcast.486.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} bitcast(%param_6.369), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.479 = f32[]{:T(128)} parameter(5) - %div.948.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%param_5.479), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.947.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} divide(%bitcast.486.clone.1, %div.948.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.311.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} select(%select_n.312.clone.1, %bitcast.486.clone.1, %div.947.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %constant.1158.clone.1 = f32[]{:T(128)} constant(0.1) - %broadcast.898.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%constant.1158.clone.1), dimensions={}, metadata={op_name="broadcast.85"} - %mul.2019.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%select_n.311.clone.1, %broadcast.898.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.129 = f32[2048,4,6144]{2,1,0:T(4,128)} parameter(8) - %constant.1162.clone.1 = f32[]{:T(128)} constant(0.9) - %broadcast.897.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%constant.1162.clone.1), dimensions={}, metadata={op_name="broadcast.84"} - %mul.2018.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%param_8.129, %broadcast.897.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.999.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} add(%mul.2019.clone.1, %mul.2018.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1325 = f32[]{:T(128)S(6)} parameter(2) - %div.944.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%param_2.1325), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %integer_pow.76.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%select_n.311.clone.1, %select_n.311.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} - %constant.1161.clone.1 = f32[]{:T(128)} constant(0.05) - %broadcast.896.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%constant.1161.clone.1), dimensions={}, metadata={op_name="broadcast.73"} - %mul.2017.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%integer_pow.76.clone.1, %broadcast.896.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.567 = f32[2048,4,6144]{2,1,0:T(4,128)} parameter(4) - %constant.1160.clone.1 = f32[]{:T(128)} constant(0.95) - %broadcast.895.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%constant.1160.clone.1), dimensions={}, metadata={op_name="broadcast.72"} - %mul.2016.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%param_4.567, %broadcast.895.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.998.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} add(%mul.2017.clone.1, %mul.2016.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1567 = f32[]{:T(128)S(6)} parameter(1) - %div.943.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%param_1.1567), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.942.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} divide(%add.998.clone.1, %div.943.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %sqrt.73.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} sqrt(%div.942.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} - %constant.1159.clone.1 = f32[]{:T(128)} constant(1e-08) - %broadcast.893.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} broadcast(%constant.1159.clone.1), dimensions={}, metadata={op_name="broadcast.65"} - %add.997.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} add(%sqrt.73.clone.1, %broadcast.893.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %multiply.437.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%div.944.clone.1, %add.997.clone.1), metadata={op_name="multiply.50"} - %div.941.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} divide(%add.999.clone.1, %multiply.437.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.2014.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%param_0.1379, %broadcast.898.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.996.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} add(%div.941.clone.1, %mul.2014.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.2013.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%mul.2015.clone.1, %add.996.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.995.clone.1 = f32[2048,4,6144]{2,1,0:T(4,128)} add(%param_0.1379, %mul.2013.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.227 = f32[2048,4,6144]{2,1,0:T(4,128)} multiply(%add.995.clone.1, %add.995.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1211 = f32[]{:T(128)} constant(0) - %reduce.186 = f32[]{:T(128)} reduce(%square.227, %constant.1211), dimensions={0,1,2}, to_apply=%region_60.65, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.189.clone.1 = f32[]{:T(128)} reduce(%integer_pow.76.clone.1, %constant.1211), dimensions={0,1,2}, to_apply=%region_45.50, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.147 = (f32[]{:T(128)}, f32[2048,4,6144]{2,1,0:T(4,128)}, f32[2048,4,6144]{2,1,0:T(4,128)}, f32[2048,4,6144]{2,1,0:T(4,128)}, f32[]{:T(128)}) tuple(%reduce.186, %add.995.clone.1, %add.998.clone.1, %add.999.clone.1, %reduce.189.clone.1) -} - -%region_39.44 (reduce_sum.302: f32[], reduce_sum.303: f32[]) -> f32[] { - %reduce_sum.303 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.302 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.304 = f32[]{:T(128)} add(%reduce_sum.302, %reduce_sum.303), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.324 (param_0.1382: f32[4,2048,16,128]) -> f32[] { - %param_0.1382 = f32[4,2048,16,128]{3,2,0,1:T(8,128)} parameter(0) - %bitcast.362 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} bitcast(%param_0.1382), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %square.230 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} multiply(%bitcast.362, %bitcast.362), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1214 = f32[]{:T(128)} constant(0) - ROOT %reduce.190 = f32[]{:T(128)} reduce(%square.230, %constant.1214), dimensions={0,1,2,3}, to_apply=%region_39.44, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} -} - -%region_38.43 (reduce_sum.296: f32[], reduce_sum.297: f32[]) -> f32[] { - %reduce_sum.297 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.296 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.298 = f32[]{:T(128)} add(%reduce_sum.296, %reduce_sum.297), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.326 (param_0.1383: f32[4,16,128,2048]) -> f32[] { - %param_0.1383 = f32[4,16,128,2048]{3,2,0,1:T(8,128)} parameter(0) - %bitcast.366 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} bitcast(%param_0.1383), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %square.233 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} multiply(%bitcast.366, %bitcast.366), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1215 = f32[]{:T(128)} constant(0) - ROOT %reduce.191 = f32[]{:T(128)} reduce(%square.233, %constant.1215), dimensions={0,1,2,3}, to_apply=%region_38.43, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} -} - -%fused_computation.327 (param_0.950: f32[16,4,128,2048]) -> bf16[4,16,128,2048] { - %param_0.950 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} parameter(0) - %copy.193 = bf16[16,4,128,2048]{3,2,0,1:T(8,128)(2,1)} copy(%param_0.950), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'self_attention\'][\'out\'][\'kernel\']"} - ROOT %bitcast.367 = bf16[4,16,128,2048]{3,2,1,0:T(8,128)(2,1)} bitcast(%copy.193), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} -} - -%region_68.73 (reduce_sum.449: f32[], reduce_sum.450: f32[]) -> f32[] { - %reduce_sum.450 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.449 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.451 = f32[]{:T(128)} add(%reduce_sum.449, %reduce_sum.450), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_53.58 (reduce_sum.368: f32[], reduce_sum.372: f32[]) -> f32[] { - %reduce_sum.372 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.368 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.373 = f32[]{:T(128)} add(%reduce_sum.368, %reduce_sum.372), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.328 (param_0.1371: f32[2048,4,16,128], param_1.1559: f32[], param_2.1317: f32[], param_3.921: f32[], param_4.559: f32[2048,4,16,128], param_5.471: f32[], param_6.361: f32[4,2048,16,128], param_7.204: pred[], param_8.121: f32[2048,4,16,128]) -> (f32[], f32[2048,4,16,128], f32[2048,4,16,128], f32[2048,4,16,128], f32[]) { - %param_0.1371 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} parameter(0) - %param_3.921 = f32[]{:T(128)S(6)} parameter(3) - %mul.1950.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} broadcast(%param_3.921), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.204 = pred[]{:T(512)S(6)} parameter(7) - %select_n.280.clone.1 = pred[2048,4,16,128]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.204), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.361 = f32[4,2048,16,128]{3,2,0,1:T(8,128)} parameter(6) - %bitcast.470.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} bitcast(%param_6.361), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.471 = f32[]{:T(128)} parameter(5) - %div.884.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} broadcast(%param_5.471), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.883.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} divide(%bitcast.470.clone.1, %div.884.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.279.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} select(%select_n.280.clone.1, %bitcast.470.clone.1, %div.883.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %constant.1110.clone.1 = f32[]{:T(128)} constant(0.1) - %broadcast.858.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} broadcast(%constant.1110.clone.1), dimensions={}, metadata={op_name="broadcast.75"} - %mul.1956.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} multiply(%select_n.279.clone.1, %broadcast.858.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.121 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} parameter(8) - %constant.1114.clone.1 = f32[]{:T(128)} constant(0.9) - %mul.1957.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} broadcast(%constant.1114.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1955.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} multiply(%param_8.121, %mul.1957.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.957.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} add(%mul.1956.clone.1, %mul.1955.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1317 = f32[]{:T(128)S(6)} parameter(2) - %div.880.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} broadcast(%param_2.1317), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %integer_pow.68.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} multiply(%select_n.279.clone.1, %select_n.279.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} - %constant.1113.clone.1 = f32[]{:T(128)} constant(0.05) - %mul.1954.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} broadcast(%constant.1113.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1952.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} multiply(%integer_pow.68.clone.1, %mul.1954.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.559 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} parameter(4) - %constant.1112.clone.1 = f32[]{:T(128)} constant(0.95) - %mul.1953.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} broadcast(%constant.1112.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1951.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} multiply(%param_4.559, %mul.1953.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.956.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} add(%mul.1952.clone.1, %mul.1951.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1559 = f32[]{:T(128)S(6)} parameter(1) - %div.879.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} broadcast(%param_1.1559), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.878.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} divide(%add.956.clone.1, %div.879.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %sqrt.65.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} sqrt(%div.878.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} - %constant.1111.clone.1 = f32[]{:T(128)} constant(1e-08) - %add.955.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} broadcast(%constant.1111.clone.1), dimensions={}, metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %add.954.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} add(%sqrt.65.clone.1, %add.955.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %multiply.429.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} multiply(%div.880.clone.1, %add.954.clone.1), metadata={op_name="multiply.58"} - %div.877.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} divide(%add.957.clone.1, %multiply.429.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1949.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} multiply(%param_0.1371, %broadcast.858.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.953.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} add(%div.877.clone.1, %mul.1949.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1948.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} multiply(%mul.1950.clone.1, %add.953.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.952.clone.1 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} add(%param_0.1371, %mul.1948.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.234 = f32[2048,4,16,128]{3,2,1,0:T(8,128)} multiply(%add.952.clone.1, %add.952.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1203 = f32[]{:T(128)} constant(0) - %reduce.192 = f32[]{:T(128)} reduce(%square.234, %constant.1203), dimensions={0,1,2,3}, to_apply=%region_68.73, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.194.clone.1 = f32[]{:T(128)} reduce(%integer_pow.68.clone.1, %constant.1203), dimensions={0,1,2,3}, to_apply=%region_53.58, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.148 = (f32[]{:T(128)}, f32[2048,4,16,128]{3,2,1,0:T(8,128)}, f32[2048,4,16,128]{3,2,1,0:T(8,128)}, f32[2048,4,16,128]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.192, %add.952.clone.1, %add.956.clone.1, %add.957.clone.1, %reduce.194.clone.1) -} - -%region_67.72 (reduce_sum.443: f32[], reduce_sum.444: f32[]) -> f32[] { - %reduce_sum.444 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.443 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.445 = f32[]{:T(128)} add(%reduce_sum.443, %reduce_sum.444), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_52.57 (reduce_sum.365: f32[], reduce_sum.366: f32[]) -> f32[] { - %reduce_sum.366 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.365 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.367 = f32[]{:T(128)} add(%reduce_sum.365, %reduce_sum.366), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.329 (param_0.1372: f32[16,4,128,2048], param_1.1560: f32[], param_2.1318: f32[], param_3.922: f32[], param_4.560: f32[16,4,128,2048], param_5.472: f32[], param_6.362: f32[4,16,128,2048], param_7.205: pred[], param_8.122: f32[16,4,128,2048]) -> (f32[], f32[16,4,128,2048], f32[16,4,128,2048], f32[16,4,128,2048], f32[]) { - %param_0.1372 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} parameter(0) - %param_3.922 = f32[]{:T(128)S(6)} parameter(3) - %mul.1960.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} broadcast(%param_3.922), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.205 = pred[]{:T(512)S(6)} parameter(7) - %select_n.284.clone.1 = pred[16,4,128,2048]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.205), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.362 = f32[4,16,128,2048]{3,2,0,1:T(8,128)} parameter(6) - %bitcast.472.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} bitcast(%param_6.362), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.472 = f32[]{:T(128)} parameter(5) - %div.892.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} broadcast(%param_5.472), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.891.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} divide(%bitcast.472.clone.1, %div.892.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.283.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} select(%select_n.284.clone.1, %bitcast.472.clone.1, %div.891.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %constant.1116.clone.1 = f32[]{:T(128)} constant(0.1) - %broadcast.860.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} broadcast(%constant.1116.clone.1), dimensions={}, metadata={op_name="broadcast.76"} - %mul.1966.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} multiply(%select_n.283.clone.1, %broadcast.860.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.122 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} parameter(8) - %constant.1120.clone.1 = f32[]{:T(128)} constant(0.9) - %mul.1967.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} broadcast(%constant.1120.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1965.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} multiply(%param_8.122, %mul.1967.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.963.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} add(%mul.1966.clone.1, %mul.1965.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1318 = f32[]{:T(128)S(6)} parameter(2) - %div.888.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} broadcast(%param_2.1318), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %integer_pow.69.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} multiply(%select_n.283.clone.1, %select_n.283.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} - %constant.1119.clone.1 = f32[]{:T(128)} constant(0.05) - %mul.1964.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} broadcast(%constant.1119.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1962.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} multiply(%integer_pow.69.clone.1, %mul.1964.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.560 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} parameter(4) - %constant.1118.clone.1 = f32[]{:T(128)} constant(0.95) - %mul.1963.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} broadcast(%constant.1118.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.1961.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} multiply(%param_4.560, %mul.1963.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.962.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} add(%mul.1962.clone.1, %mul.1961.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1560 = f32[]{:T(128)S(6)} parameter(1) - %div.887.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} broadcast(%param_1.1560), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.886.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} divide(%add.962.clone.1, %div.887.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %sqrt.66.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} sqrt(%div.886.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} - %constant.1117.clone.1 = f32[]{:T(128)} constant(1e-08) - %add.961.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} broadcast(%constant.1117.clone.1), dimensions={}, metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %add.960.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} add(%sqrt.66.clone.1, %add.961.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %multiply.430.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} multiply(%div.888.clone.1, %add.960.clone.1), metadata={op_name="multiply.57"} - %div.885.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} divide(%add.963.clone.1, %multiply.430.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1959.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} multiply(%param_0.1372, %broadcast.860.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.959.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} add(%div.885.clone.1, %mul.1959.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1958.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} multiply(%mul.1960.clone.1, %add.959.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.958.clone.1 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} add(%param_0.1372, %mul.1958.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.235 = f32[16,4,128,2048]{3,2,1,0:T(8,128)} multiply(%add.958.clone.1, %add.958.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1204 = f32[]{:T(128)} constant(0) - %reduce.193 = f32[]{:T(128)} reduce(%square.235, %constant.1204), dimensions={0,1,2,3}, to_apply=%region_67.72, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.195.clone.1 = f32[]{:T(128)} reduce(%integer_pow.69.clone.1, %constant.1204), dimensions={0,1,2,3}, to_apply=%region_52.57, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.149 = (f32[]{:T(128)}, f32[16,4,128,2048]{3,2,1,0:T(8,128)}, f32[16,4,128,2048]{3,2,1,0:T(8,128)}, f32[16,4,128,2048]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.193, %add.958.clone.1, %add.962.clone.1, %add.963.clone.1, %reduce.195.clone.1) -} - -%region_41.46 (reduce_sum.311: f32[], reduce_sum.312: f32[]) -> f32[] { - %reduce_sum.312 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.311 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.316 = f32[]{:T(128)} add(%reduce_sum.311, %reduce_sum.312), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_36.41 (reduce_sum.284: f32[], reduce_sum.288: f32[]) -> f32[] { - %reduce_sum.288 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.284 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.289 = f32[]{:T(128)} add(%reduce_sum.284, %reduce_sum.288), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.341 (param_0.1385: f32[4,2048,8,128], param_1.1571: f32[4,2048,8,128]) -> (f32[], f32[]) { - %param_0.1385 = f32[4,2048,8,128]{3,2,0,1:T(8,128)S(1)} parameter(0) - %bitcast.371 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_0.1385), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %square.238 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%bitcast.371, %bitcast.371), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1217 = f32[]{:T(128)} constant(0) - %reduce.196 = f32[]{:T(128)} reduce(%square.238, %constant.1217), dimensions={0,1,2,3}, to_apply=%region_41.46, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %param_1.1571 = f32[4,2048,8,128]{3,2,0,1:T(8,128)} parameter(1) - %bitcast.375.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_1.1571), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %square.241.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%bitcast.375.clone.1, %bitcast.375.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %reduce.197.clone.1 = f32[]{:T(128)} reduce(%square.241.clone.1, %constant.1217), dimensions={0,1,2,3}, to_apply=%region_36.41, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.168 = (f32[]{:T(128)}, f32[]{:T(128)}) tuple(%reduce.196, %reduce.197.clone.1) -} - -%fused_computation.344 (param_0.982: f32[2048,4,8,128]) -> bf16[4,2048,8,128] { - %param_0.982 = f32[2048,4,8,128]{3,2,1,0:T(8,128)S(1)} parameter(0) - %copy.194 = bf16[2048,4,8,128]{3,2,0,1:T(8,128)(2,1)} copy(%param_0.982), sharding={replicated}, metadata={op_name="state.params[\'params\'][\'decoder\'][\'layers\'][\'self_attention\'][\'value\'][\'kernel\']"} - ROOT %bitcast.376 = bf16[4,2048,8,128]{3,2,1,0:T(8,128)(2,1)} bitcast(%copy.194), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} -} - -%region_70.75 (reduce_sum.458: f32[], reduce_sum.459: f32[]) -> f32[] { - %reduce_sum.459 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.458 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.463 = f32[]{:T(128)} add(%reduce_sum.458, %reduce_sum.459), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_55.60 (reduce_sum.380: f32[], reduce_sum.381: f32[]) -> f32[] { - %reduce_sum.381 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.380 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.382 = f32[]{:T(128)} add(%reduce_sum.380, %reduce_sum.381), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.345 (param_0.1369: f32[2048,4,8,128], param_1.1557: f32[], param_2.1315: f32[], param_3.919: f32[], param_4.557: f32[2048,4,8,128], param_5.469: f32[], param_6.359: f32[4,2048,8,128], param_7.202: pred[], param_8.119: f32[2048,4,8,128]) -> (f32[], f32[2048,4,8,128], f32[2048,4,8,128], f32[2048,4,8,128], f32[]) { - %param_0.1369 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} parameter(0) - %param_3.919 = f32[]{:T(128)S(6)} parameter(3) - %mul.1936.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_3.919), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.202 = pred[]{:T(512)S(6)} parameter(7) - %select_n.272.clone.1 = pred[2048,4,8,128]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.202), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.359 = f32[4,2048,8,128]{3,2,0,1:T(8,128)} parameter(6) - %bitcast.466.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_6.359), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.469 = f32[]{:T(128)} parameter(5) - %div.868.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_5.469), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.867.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} divide(%bitcast.466.clone.1, %div.868.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.271.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} select(%select_n.272.clone.1, %bitcast.466.clone.1, %div.867.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %constant.1098.clone.1 = f32[]{:T(128)} constant(0.1) - %broadcast.850.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.1098.clone.1), dimensions={}, metadata={op_name="broadcast.80"} - %mul.1940.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%select_n.271.clone.1, %broadcast.850.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.119 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} parameter(8) - %constant.1102.clone.1 = f32[]{:T(128)} constant(0.9) - %broadcast.849.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.1102.clone.1), dimensions={}, metadata={op_name="broadcast.79"} - %mul.1939.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_8.119, %broadcast.849.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.946.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} add(%mul.1940.clone.1, %mul.1939.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1315 = f32[]{:T(128)S(6)} parameter(2) - %div.864.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_2.1315), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %integer_pow.66.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%select_n.271.clone.1, %select_n.271.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} - %constant.1101.clone.1 = f32[]{:T(128)} constant(0.05) - %broadcast.848.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.1101.clone.1), dimensions={}, metadata={op_name="broadcast.69"} - %mul.1938.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%integer_pow.66.clone.1, %broadcast.848.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.557 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} parameter(4) - %constant.1100.clone.1 = f32[]{:T(128)} constant(0.95) - %broadcast.847.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.1100.clone.1), dimensions={}, metadata={op_name="broadcast.68"} - %mul.1937.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_4.557, %broadcast.847.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.945.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} add(%mul.1938.clone.1, %mul.1937.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1557 = f32[]{:T(128)S(6)} parameter(1) - %div.863.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_1.1557), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.862.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} divide(%add.945.clone.1, %div.863.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %sqrt.63.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} sqrt(%div.862.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} - %constant.1099.clone.1 = f32[]{:T(128)} constant(1e-08) - %broadcast.845.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.1099.clone.1), dimensions={}, metadata={op_name="broadcast.63"} - %add.944.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} add(%sqrt.63.clone.1, %broadcast.845.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %multiply.427.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%div.864.clone.1, %add.944.clone.1), metadata={op_name="multiply.60"} - %div.861.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} divide(%add.946.clone.1, %multiply.427.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1935.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_0.1369, %broadcast.850.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.943.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} add(%div.861.clone.1, %mul.1935.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1934.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%mul.1936.clone.1, %add.943.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.942.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} add(%param_0.1369, %mul.1934.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.242 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%add.942.clone.1, %add.942.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1201 = f32[]{:T(128)} constant(0) - %reduce.198 = f32[]{:T(128)} reduce(%square.242, %constant.1201), dimensions={0,1,2,3}, to_apply=%region_70.75, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.200.clone.1 = f32[]{:T(128)} reduce(%integer_pow.66.clone.1, %constant.1201), dimensions={0,1,2,3}, to_apply=%region_55.60, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.150 = (f32[]{:T(128)}, f32[2048,4,8,128]{3,2,1,0:T(8,128)}, f32[2048,4,8,128]{3,2,1,0:T(8,128)}, f32[2048,4,8,128]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.198, %add.942.clone.1, %add.945.clone.1, %add.946.clone.1, %reduce.200.clone.1) -} - -%region_65.70 (reduce_sum.431: f32[], reduce_sum.435: f32[]) -> f32[] { - %reduce_sum.435 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.431 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.436 = f32[]{:T(128)} add(%reduce_sum.431, %reduce_sum.435), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_50.55 (reduce_sum.353: f32[], reduce_sum.354: f32[]) -> f32[] { - %reduce_sum.354 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.353 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.358 = f32[]{:T(128)} add(%reduce_sum.353, %reduce_sum.354), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.346 (param_0.1374: f32[2048,4,8,128], param_1.1562: f32[], param_2.1320: f32[], param_3.924: f32[], param_4.562: f32[2048,4,8,128], param_5.474: f32[], param_6.364: f32[4,2048,8,128], param_7.207: pred[], param_8.124: f32[2048,4,8,128]) -> (f32[], f32[2048,4,8,128], f32[2048,4,8,128], f32[2048,4,8,128], f32[]) { - %param_0.1374 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} parameter(0) - %param_3.924 = f32[]{:T(128)S(6)} parameter(3) - %mul.1977.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_3.924), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.207 = pred[]{:T(512)S(6)} parameter(7) - %select_n.292.clone.1 = pred[2048,4,8,128]{3,2,1,0:T(8,128)(4,1)} broadcast(%param_7.207), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.364 = f32[4,2048,8,128]{3,2,0,1:T(8,128)S(1)} parameter(6) - %bitcast.476.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} bitcast(%param_6.364), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.474 = f32[]{:T(128)} parameter(5) - %div.908.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_5.474), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.907.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} divide(%bitcast.476.clone.1, %div.908.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.291.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} select(%select_n.292.clone.1, %bitcast.476.clone.1, %div.907.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %constant.1128.clone.1 = f32[]{:T(128)} constant(0.1) - %broadcast.872.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.1128.clone.1), dimensions={}, metadata={op_name="broadcast.80"} - %mul.1981.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%select_n.291.clone.1, %broadcast.872.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.124 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} parameter(8) - %constant.1132.clone.1 = f32[]{:T(128)} constant(0.9) - %broadcast.871.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.1132.clone.1), dimensions={}, metadata={op_name="broadcast.79"} - %mul.1980.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_8.124, %broadcast.871.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.973.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} add(%mul.1981.clone.1, %mul.1980.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1320 = f32[]{:T(128)S(6)} parameter(2) - %div.904.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_2.1320), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %integer_pow.71.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%select_n.291.clone.1, %select_n.291.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} - %constant.1131.clone.1 = f32[]{:T(128)} constant(0.05) - %broadcast.870.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.1131.clone.1), dimensions={}, metadata={op_name="broadcast.69"} - %mul.1979.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%integer_pow.71.clone.1, %broadcast.870.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.562 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} parameter(4) - %constant.1130.clone.1 = f32[]{:T(128)} constant(0.95) - %broadcast.869.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.1130.clone.1), dimensions={}, metadata={op_name="broadcast.68"} - %mul.1978.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_4.562, %broadcast.869.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.972.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} add(%mul.1979.clone.1, %mul.1978.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1562 = f32[]{:T(128)S(6)} parameter(1) - %div.903.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%param_1.1562), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.902.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} divide(%add.972.clone.1, %div.903.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %sqrt.68.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} sqrt(%div.902.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} - %constant.1129.clone.1 = f32[]{:T(128)} constant(1e-08) - %broadcast.867.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} broadcast(%constant.1129.clone.1), dimensions={}, metadata={op_name="broadcast.63"} - %add.971.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} add(%sqrt.68.clone.1, %broadcast.867.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %multiply.432.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%div.904.clone.1, %add.971.clone.1), metadata={op_name="multiply.55"} - %div.901.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} divide(%add.973.clone.1, %multiply.432.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1976.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%param_0.1374, %broadcast.872.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.970.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} add(%div.901.clone.1, %mul.1976.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1975.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%mul.1977.clone.1, %add.970.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.969.clone.1 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} add(%param_0.1374, %mul.1975.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.243 = f32[2048,4,8,128]{3,2,1,0:T(8,128)} multiply(%add.969.clone.1, %add.969.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1206 = f32[]{:T(128)} constant(0) - %reduce.199 = f32[]{:T(128)} reduce(%square.243, %constant.1206), dimensions={0,1,2,3}, to_apply=%region_65.70, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.201.clone.1 = f32[]{:T(128)} reduce(%integer_pow.71.clone.1, %constant.1206), dimensions={0,1,2,3}, to_apply=%region_50.55, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.151 = (f32[]{:T(128)}, f32[2048,4,8,128]{3,2,1,0:T(8,128)}, f32[2048,4,8,128]{3,2,1,0:T(8,128)}, f32[2048,4,8,128]{3,2,1,0:T(8,128)}, f32[]{:T(128)}) tuple(%reduce.199, %add.969.clone.1, %add.972.clone.1, %add.973.clone.1, %reduce.201.clone.1) -} - -%fused_computation.362 (param_0.1056: bf16[4,128,2048], param_1.1117: f32[4,128], param_2.830: f32[4,128], param_3.495: bf16[4,128,2048], param_4.296: bf16[2048]) -> bf16[4,128,2048] { - %param_3.495 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) - %param_4.296 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(4) - %dot_general.448 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_4.296), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %dot_general.438 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%param_3.495, %dot_general.448), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %convert_element_type.1363 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%dot_general.438), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %param_2.830 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %mul.1851 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_2.830), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %mul.1843 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1363, %mul.1851), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %param_0.1056 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) - %convert_element_type.1374 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%param_0.1056), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %param_1.1117 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %mul.1850 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_1.1117), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %mul.1849 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1374, %mul.1850), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %add_any.193 = f32[4,128,2048]{2,1,0:T(8,128)} add(%mul.1843, %mul.1849), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/add_any" stack_frame_id=0} - ROOT %convert_element_type.1361 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} convert(%add_any.193), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} +%fused_computation.491 (param_0.1515: f32[2048,4]) -> bf16[4,2048] { + %param_0.1515 = f32[2048,4]{0,1:T(4,128)S(1)} parameter(0) + %bitcast.755 = f32[4,2048]{1,0:T(4,128)} bitcast(%param_0.1515), metadata={op_name="jit(train_step)/jvp()/transpose" stack_frame_id=0} + ROOT %convert.172 = bf16[4,2048]{1,0:T(4,128)(2,1)S(1)} convert(%bitcast.755) } -%region_7.10 (reduce_sum.171: f32[], reduce_sum.184: f32[]) -> f32[] { - %reduce_sum.184 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/reduce_sum"} - %reduce_sum.171 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/reduce_sum"} - ROOT %reduce_sum.185 = f32[]{:T(128)} add(%reduce_sum.171, %reduce_sum.184), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +%fused_computation.486 (param_0.1352: s32[1,128]) -> (f32[1,128,1,1], f32[128]) { + %param_0.1352 = s32[1,128]{1,0:T(1,128)S(1)} parameter(0) + %convert_element_type.1325 = f32[1,128]{1,0:T(1,128)} convert(%param_0.1352), metadata={op_name="jit(train_step)/convert_element_type" stack_frame_id=0} + %bitcast.626 = f32[1,128,1,1]{1,3,2,0:T(1,128)} bitcast(%convert_element_type.1325), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %bitcast.627.clone.1 = f32[128]{0:T(128)S(1)} bitcast(%convert_element_type.1325), metadata={op_name="jit(train_step)/convert_element_type" stack_frame_id=0} + ROOT %tuple.170 = (f32[1,128,1,1]{1,3,2,0:T(1,128)}, f32[128]{0:T(128)S(1)}) tuple(%bitcast.626, %bitcast.627.clone.1) +} + +%fused_computation.485 () -> f32[64] { + %constant.466.clone.1 = f32[]{:T(128)} constant(1e+06) + %broadcast.975 = f32[64]{0:T(128)} broadcast(%constant.466.clone.1), dimensions={}, metadata={op_name="broadcast.409"} + %iota.56 = s32[64]{0:T(128)} iota(), iota_dimension=0, metadata={op_name="jit(train_step)/iota" stack_frame_id=0} + %constant.467.clone.1 = s32[]{:T(128)} constant(2) + %broadcast.974 = s32[64]{0:T(128)} broadcast(%constant.467.clone.1), dimensions={}, metadata={op_name="broadcast.410"} + %mul.1960 = s32[64]{0:T(128)} multiply(%iota.56, %broadcast.974), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} + %convert_element_type.1324 = f32[64]{0:T(128)} convert(%mul.1960), metadata={op_name="jit(train_step)/convert_element_type" stack_frame_id=0} + %constant.468.clone.1 = f32[]{:T(128)} constant(0.0078125) + %broadcast.973 = f32[64]{0:T(128)} broadcast(%constant.468.clone.1), dimensions={}, metadata={op_name="broadcast.411"} + %div.778 = f32[64]{0:T(128)} multiply(%convert_element_type.1324, %broadcast.973), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + ROOT %pow.36 = f32[64]{0:T(128)S(1)} power(%broadcast.975, %div.778), metadata={op_name="jit(train_step)/pow" stack_frame_id=0} } -%fused_computation.363 (param_0.1394: bf16[4,128,2048]) -> f32[4,128] { - %param_0.1394 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) - %convert_element_type.1365 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%param_0.1394), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %square.246 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1365, %convert_element_type.1365), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/square" stack_frame_id=0} - %constant.1227 = f32[]{:T(128)} constant(0) - ROOT %reduce.202 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%square.246, %constant.1227), dimensions={2}, to_apply=%region_7.10, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/reduce_sum" stack_frame_id=0} +%fused_computation.421 (param_0.1238: f32[128], param_1.1308: f32[64]) -> (bf16[1,128,1,64], bf16[1,128,1,64]) { + %param_0.1238 = f32[128]{0:T(128)S(1)} parameter(0) + %div.732 = f32[1,128,1,64]{3,1,2,0:T(8,128)} broadcast(%param_0.1238), dimensions={1}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %param_1.1308 = f32[64]{0:T(128)S(1)} parameter(1) + %div.738 = f32[1,128,1,64]{3,1,2,0:T(8,128)} broadcast(%param_1.1308), dimensions={3}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %div.729 = f32[1,128,1,64]{3,1,2,0:T(8,128)} divide(%div.732, %div.738), metadata={op_name="jit(train_step)/div" stack_frame_id=0} + %sin.38 = f32[1,128,1,64]{3,1,2,0:T(8,128)} sine(%div.729), metadata={op_name="jit(train_step)/sin" stack_frame_id=0} + %convert_element_type.1310 = bf16[1,128,1,64]{3,1,2,0:T(8,128)(2,1)S(1)} convert(%sin.38), metadata={op_name="jit(train_step)/convert_element_type" stack_frame_id=0} + %cos.41.clone.1 = f32[1,128,1,64]{3,1,2,0:T(8,128)} cosine(%div.729), metadata={op_name="jit(train_step)/cos" stack_frame_id=0} + %convert_element_type.1309.clone.1 = bf16[1,128,1,64]{3,1,2,0:T(8,128)(2,1)S(1)} convert(%cos.41.clone.1), metadata={op_name="jit(train_step)/convert_element_type" stack_frame_id=0} + ROOT %tuple.167 = (bf16[1,128,1,64]{3,1,2,0:T(8,128)(2,1)S(1)}, bf16[1,128,1,64]{3,1,2,0:T(8,128)(2,1)S(1)}) tuple(%convert_element_type.1310, %convert_element_type.1309.clone.1) } -%region_12.15 (reduce_sum.198: f32[], reduce_sum.199: f32[]) -> f32[] { - %reduce_sum.199 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/reduce_sum"} - %reduce_sum.198 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/reduce_sum"} - ROOT %reduce_sum.200 = f32[]{:T(128)} add(%reduce_sum.198, %reduce_sum.199), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.365 (param_0.1389: bf16[4,128,2048], param_1.1574: bf16[4,128,2048], param_2.1328: bf16[2048]) -> f32[4,128] { - %param_0.1389 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) - %convert_element_type.1372 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%param_0.1389), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %param_1.1574 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %param_2.1328 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(2) - %dot_general.447 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1328), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %dot_general.437 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%param_1.1574, %dot_general.447), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %convert_element_type.1371 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%dot_general.437), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %mul.1847 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1372, %convert_element_type.1371), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %constant.1221 = f32[]{:T(128)} constant(0) - ROOT %reduce.203 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%mul.1847, %constant.1221), dimensions={2}, to_apply=%region_12.15, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/reduce_sum" stack_frame_id=0} -} - -%region_10.13 (dot_general.190: bf16[], dot_general.191: bf16[]) -> bf16[] { - %dot_general.191 = bf16[]{:T(256)} parameter(1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general"} - %dot_general.190 = bf16[]{:T(256)} parameter(0), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general"} - ROOT %add.419 = bf16[]{:T(256)} add(%dot_general.190, %dot_general.191), metadata={op_name="add.82"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.285.clone.clone (param_0.1351: bf16[151936,2048]) -> bf16[151936,2048,1] { - %param_0.1351 = bf16[151936,2048]{1,0:T(8,128)(2,1)} parameter(0) - ROOT %bitcast.528 = bf16[151936,2048,1]{1,0,2:T(8,128)(2,1)} bitcast(%param_0.1351), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} -} - -%fused_computation.289.clone.1.clone.clone (param_0.1352: bf16[4,128,151936], param_1.1546: s32[4,128], param_2.1285: f32[4,128], param_3.906: f32[4,128], param_4.542: bf16[4,128], param_5.442: f32[4,128]) -> bf16[4,128,151936] { - %param_5.442 = f32[4,128]{1,0:T(4,128)S(1)} parameter(5) - %mul.2075 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_5.442), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - %param_3.906 = f32[4,128]{1,0:T(4,128)S(1)} parameter(3) - %mul.2074 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_3.906), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - %param_0.1352 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} parameter(0) - %convert_element_type.1444 = f32[4,128,151936]{2,1,0:T(8,128)} convert(%param_0.1352), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %param_4.542 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(4) - %sub.88 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_4.542), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %sub.87 = f32[4,128,151936]{2,1,0:T(8,128)} subtract(%convert_element_type.1444, %sub.88), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} - %exp.60 = f32[4,128,151936]{2,1,0:T(8,128)} exponential(%sub.87), metadata={op_name="jit(train_step)/jvp()/exp" stack_frame_id=0} - %mul.2073 = f32[4,128,151936]{2,1,0:T(8,128)} multiply(%mul.2074, %exp.60), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - %param_2.1285 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %div.962 = f32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_2.1285), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} - %div.961 = f32[4,128,151936]{2,1,0:T(8,128)} divide(%mul.2073, %div.962), metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} - %param_1.1546 = s32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %eq.43 = s32[4,128,151936]{2,1,0:T(8,128)} broadcast(%param_1.1546), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} - %eq.42 = s32[4,128,151936]{2,1,0:T(8,128)} iota(), iota_dimension=2, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} - %eq.41 = pred[4,128,151936]{2,1,0:T(8,128)(4,1)} compare(%eq.43, %eq.42), direction=EQ, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} - %convert_element_type.1443 = f32[4,128,151936]{2,1,0:T(8,128)} convert(%eq.41), metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/convert_element_type" stack_frame_id=0} - %sub.86 = f32[4,128,151936]{2,1,0:T(8,128)} subtract(%div.961, %convert_element_type.1443), metadata={op_name="jit(train_step)/transpose(jvp())/sub" stack_frame_id=0} - %mul.2072 = f32[4,128,151936]{2,1,0:T(8,128)} multiply(%mul.2075, %sub.86), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - ROOT %convert_element_type.1442 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} convert(%mul.2072), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} -} - -%fused_computation.366 (param_0.1350: f32[4,128], param_1.1545: bf16[4,128,2048], param_2.1286: bf16[151936,2048], param_3.907: bf16[4,128,151936], param_4.543: s32[4,128], param_5.443: f32[4,128], param_6.340: f32[4,128], param_7.199: bf16[4,128], param_8.116: f32[4,128]) -> (bf16[2048], bf16[4,128,2048]) { - %param_3.907 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} parameter(3) - %param_4.543 = s32[4,128]{1,0:T(4,128)S(1)} parameter(4) - %param_5.443 = f32[4,128]{1,0:T(4,128)S(1)} parameter(5) - %param_6.340 = f32[4,128]{1,0:T(4,128)S(1)} parameter(6) - %param_7.199 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(7) - %param_8.116 = f32[4,128]{1,0:T(4,128)S(1)} parameter(8) - %multiply_convert_fusion.2.clone.1 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} fusion(%param_3.907, %param_4.543, %param_5.443, %param_6.340, %param_7.199, /*index=5*/%param_8.116), kind=kLoop, calls=%fused_computation.289.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/convert_element_type" stack_frame_id=0} - %param_2.1286 = bf16[151936,2048]{1,0:T(8,128)(2,1)} parameter(2) - %fusion.251.clone.1 = bf16[151936,2048,1]{1,0,2:T(8,128)(2,1)} fusion(%param_2.1286), kind=kLoop, calls=%fused_computation.285.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} - %convolution.84.clone.1 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} convolution(%multiply_convert_fusion.2.clone.1, %fusion.251.clone.1), window={size=1}, dim_labels=0bf_io0->0bf, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/dot_general" stack_frame_id=0} - %param_1.1545 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %convert_element_type.1384 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%param_1.1545), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %param_0.1350 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.1862 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_0.1350), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %mul.1861 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1384, %mul.1862), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %convert_element_type.1383 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.1861), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %multiply.420 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%convolution.84.clone.1, %convert_element_type.1383), metadata={op_name="multiply.362"} - %constant.1050 = bf16[]{:T(256)} constant(0) - %reduce.204 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} reduce(%multiply.420, %constant.1050), dimensions={0,1}, to_apply=%region_10.13, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - ROOT %tuple.165 = (bf16[2048]{0:T(1024)(128)(2,1)S(1)}, bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)}) tuple(%reduce.204, %convolution.84.clone.1) -} - -%fused_computation.374 (param_0.1088: f32[64], param_1.1150: f32[4,128]) -> (bf16[4,128,1,64], bf16[4,128,1,64]) { - %param_1.1150 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %div.720 = f32[4,128,1,64]{3,1,0,2:T(8,128)} broadcast(%param_1.1150), dimensions={0,1}, metadata={op_name="jit(train_step)/layers/div" stack_frame_id=0} - %param_0.1088 = f32[64]{0:T(128)S(1)} parameter(0) - %div.718 = f32[4,128,1,64]{3,1,0,2:T(8,128)} broadcast(%param_0.1088), dimensions={3}, metadata={op_name="jit(train_step)/layers/div" stack_frame_id=0} - %div.717 = f32[4,128,1,64]{3,1,0,2:T(8,128)} divide(%div.720, %div.718), metadata={op_name="jit(train_step)/layers/div" stack_frame_id=0} - %sin.38 = f32[4,128,1,64]{3,1,0,2:T(8,128)} sine(%div.717), metadata={op_name="jit(train_step)/layers/sin" stack_frame_id=0} - %convert_element_type.1392 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} convert(%sin.38), metadata={op_name="jit(train_step)/layers/convert_element_type" stack_frame_id=0} - %cos.41.clone.1 = f32[4,128,1,64]{3,1,0,2:T(8,128)} cosine(%div.717), metadata={op_name="jit(train_step)/layers/cos" stack_frame_id=0} - %convert_element_type.1391.clone.1 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} convert(%cos.41.clone.1), metadata={op_name="jit(train_step)/layers/convert_element_type" stack_frame_id=0} - ROOT %tuple.158 = (bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)}, bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)}) tuple(%convert_element_type.1392, %convert_element_type.1391.clone.1) -} - -%fused_computation.375 (param_0.1085: bf16[4,128,1,64]) -> bf16[4,128,1,128] { - %param_0.1085 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} parameter(0) - %constant.1042 = bf16[]{:T(256)} constant(-inf) - %pad.46 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1085, %constant.1042), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} - %pad.45 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1085, %constant.1042), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} - ROOT %maximum.42 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} maximum(%pad.46, %pad.45), metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} -} - -%fused_computation.376 (param_0.1087: bf16[4,128,1,64]) -> bf16[4,128,1,128] { - %param_0.1087 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} parameter(0) - %constant.1041 = bf16[]{:T(256)} constant(-inf) - %pad.48 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1087, %constant.1041), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} - %pad.47 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1087, %constant.1041), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} - ROOT %maximum.43 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} maximum(%pad.48, %pad.47), metadata={op_name="jit(train_step)/layers/concatenate" stack_frame_id=0} -} - -%region_35.40 (reduce_sum.281: f32[], reduce_sum.282: f32[]) -> f32[] { - %reduce_sum.282 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.281 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.283 = f32[]{:T(128)} add(%reduce_sum.281, %reduce_sum.282), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_34.39 (reduce_sum.275: f32[], reduce_sum.276: f32[]) -> f32[] { - %reduce_sum.276 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.275 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.277 = f32[]{:T(128)} add(%reduce_sum.275, %reduce_sum.276), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.380 (param_0.1386: f32[4,2048], param_1.1572: f32[4,2048]) -> (f32[], f32[]) { - %param_0.1386 = f32[4,2048]{1,0:T(4,128)S(1)} parameter(0) - %bitcast.404 = f32[2048,4]{0,1:T(4,128)} bitcast(%param_0.1386), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %square.249 = f32[2048,4]{0,1:T(4,128)} multiply(%bitcast.404, %bitcast.404), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1218 = f32[]{:T(128)} constant(0) - %reduce.205 = f32[]{:T(128)} reduce(%square.249, %constant.1218), dimensions={0,1}, to_apply=%region_35.40, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %param_1.1572 = f32[4,2048]{1,0:T(4,128)} parameter(1) - %bitcast.408.clone.1 = f32[2048,4]{0,1:T(4,128)} bitcast(%param_1.1572), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %square.252.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%bitcast.408.clone.1, %bitcast.408.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %reduce.206.clone.1 = f32[]{:T(128)} reduce(%square.252.clone.1, %constant.1218), dimensions={0,1}, to_apply=%region_34.39, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.169 = (f32[]{:T(128)}, f32[]{:T(128)}) tuple(%reduce.205, %reduce.206.clone.1) -} - -%region_64.69 (reduce_sum.428: f32[], reduce_sum.429: f32[]) -> f32[] { - %reduce_sum.429 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.428 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.430 = f32[]{:T(128)} add(%reduce_sum.428, %reduce_sum.429), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_49.54 (reduce_sum.347: f32[], reduce_sum.351: f32[]) -> f32[] { - %reduce_sum.351 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.347 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.352 = f32[]{:T(128)} add(%reduce_sum.347, %reduce_sum.351), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.383 (param_0.1375: f32[2048,4], param_1.1563: f32[], param_2.1321: f32[], param_3.925: f32[], param_4.563: f32[2048,4], param_5.475: f32[], param_6.365: f32[4,2048], param_7.208: pred[], param_8.125: f32[2048,4]) -> (f32[], f32[2048,4], f32[2048,4], f32[2048,4], f32[]) { - %param_0.1375 = f32[2048,4]{0,1:T(4,128)S(1)} parameter(0) - %param_3.925 = f32[]{:T(128)S(6)} parameter(3) - %mul.1984.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_3.925), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.208 = pred[]{:T(512)S(6)} parameter(7) - %select_n.296.clone.1 = pred[2048,4]{0,1:T(4,128)(4,1)} broadcast(%param_7.208), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.365 = f32[4,2048]{1,0:T(4,128)S(1)} parameter(6) - %bitcast.478.clone.1 = f32[2048,4]{0,1:T(4,128)} bitcast(%param_6.365), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.475 = f32[]{:T(128)} parameter(5) - %div.916.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_5.475), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.915.clone.1 = f32[2048,4]{0,1:T(4,128)} divide(%bitcast.478.clone.1, %div.916.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.295.clone.1 = f32[2048,4]{0,1:T(4,128)} select(%select_n.296.clone.1, %bitcast.478.clone.1, %div.915.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %constant.1134.clone.1 = f32[]{:T(128)} constant(0.1) - %broadcast.878.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%constant.1134.clone.1), dimensions={}, metadata={op_name="broadcast.82"} - %mul.1988.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%select_n.295.clone.1, %broadcast.878.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.125 = f32[2048,4]{0,1:T(4,128)S(1)} parameter(8) - %constant.1138.clone.1 = f32[]{:T(128)} constant(0.9) - %broadcast.877.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%constant.1138.clone.1), dimensions={}, metadata={op_name="broadcast.81"} - %mul.1987.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%param_8.125, %broadcast.877.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.978.clone.1 = f32[2048,4]{0,1:T(4,128)S(1)} add(%mul.1988.clone.1, %mul.1987.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1321 = f32[]{:T(128)S(6)} parameter(2) - %div.912.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_2.1321), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %integer_pow.72.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%select_n.295.clone.1, %select_n.295.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} - %constant.1137.clone.1 = f32[]{:T(128)} constant(0.05) - %broadcast.876.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%constant.1137.clone.1), dimensions={}, metadata={op_name="broadcast.71"} - %mul.1986.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%integer_pow.72.clone.1, %broadcast.876.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.563 = f32[2048,4]{0,1:T(4,128)S(1)} parameter(4) - %constant.1136.clone.1 = f32[]{:T(128)} constant(0.95) - %broadcast.875.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%constant.1136.clone.1), dimensions={}, metadata={op_name="broadcast.70"} - %mul.1985.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%param_4.563, %broadcast.875.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.977.clone.1 = f32[2048,4]{0,1:T(4,128)S(1)} add(%mul.1986.clone.1, %mul.1985.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1563 = f32[]{:T(128)S(6)} parameter(1) - %div.911.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_1.1563), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.910.clone.1 = f32[2048,4]{0,1:T(4,128)} divide(%add.977.clone.1, %div.911.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %sqrt.69.clone.1 = f32[2048,4]{0,1:T(4,128)} sqrt(%div.910.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} - %constant.1135.clone.1 = f32[]{:T(128)} constant(1e-08) - %broadcast.873.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%constant.1135.clone.1), dimensions={}, metadata={op_name="broadcast.64"} - %add.976.clone.1 = f32[2048,4]{0,1:T(4,128)} add(%sqrt.69.clone.1, %broadcast.873.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %multiply.433.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%div.912.clone.1, %add.976.clone.1), metadata={op_name="multiply.54"} - %div.909.clone.1 = f32[2048,4]{0,1:T(4,128)} divide(%add.978.clone.1, %multiply.433.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1983.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%param_0.1375, %broadcast.878.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.975.clone.1 = f32[2048,4]{0,1:T(4,128)} add(%div.909.clone.1, %mul.1983.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1982.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%mul.1984.clone.1, %add.975.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.974.clone.1 = f32[2048,4]{0,1:T(4,128)S(1)} add(%param_0.1375, %mul.1982.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.253 = f32[2048,4]{0,1:T(4,128)} multiply(%add.974.clone.1, %add.974.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1207 = f32[]{:T(128)} constant(0) - %reduce.207 = f32[]{:T(128)} reduce(%square.253, %constant.1207), dimensions={0,1}, to_apply=%region_64.69, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.209.clone.1 = f32[]{:T(128)} reduce(%integer_pow.72.clone.1, %constant.1207), dimensions={0,1}, to_apply=%region_49.54, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.152 = (f32[]{:T(128)}, f32[2048,4]{0,1:T(4,128)S(1)}, f32[2048,4]{0,1:T(4,128)S(1)}, f32[2048,4]{0,1:T(4,128)S(1)}, f32[]{:T(128)}) tuple(%reduce.207, %add.974.clone.1, %add.977.clone.1, %add.978.clone.1, %reduce.209.clone.1) -} - -%region_63.68 (reduce_sum.422: f32[], reduce_sum.423: f32[]) -> f32[] { - %reduce_sum.423 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.422 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.424 = f32[]{:T(128)} add(%reduce_sum.422, %reduce_sum.423), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_48.53 (reduce_sum.344: f32[], reduce_sum.345: f32[]) -> f32[] { - %reduce_sum.345 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.344 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.346 = f32[]{:T(128)} add(%reduce_sum.344, %reduce_sum.345), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.384 (param_0.1376: f32[2048,4], param_1.1564: f32[], param_2.1322: f32[], param_3.926: f32[], param_4.564: f32[2048,4], param_5.476: f32[], param_6.366: f32[4,2048], param_7.209: pred[], param_8.126: f32[2048,4]) -> (f32[], f32[2048,4], f32[2048,4], f32[2048,4], f32[]) { - %param_0.1376 = f32[2048,4]{0,1:T(4,128)S(1)} parameter(0) - %param_3.926 = f32[]{:T(128)S(6)} parameter(3) - %mul.1991.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_3.926), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.209 = pred[]{:T(512)S(6)} parameter(7) - %select_n.300.clone.1 = pred[2048,4]{0,1:T(4,128)(4,1)} broadcast(%param_7.209), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.366 = f32[4,2048]{1,0:T(4,128)} parameter(6) - %bitcast.480.clone.1 = f32[2048,4]{0,1:T(4,128)} bitcast(%param_6.366), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.476 = f32[]{:T(128)} parameter(5) - %div.924.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_5.476), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.923.clone.1 = f32[2048,4]{0,1:T(4,128)} divide(%bitcast.480.clone.1, %div.924.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.299.clone.1 = f32[2048,4]{0,1:T(4,128)} select(%select_n.300.clone.1, %bitcast.480.clone.1, %div.923.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %constant.1140.clone.1 = f32[]{:T(128)} constant(0.1) - %broadcast.884.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%constant.1140.clone.1), dimensions={}, metadata={op_name="broadcast.82"} - %mul.1995.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%select_n.299.clone.1, %broadcast.884.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.126 = f32[2048,4]{0,1:T(4,128)S(1)} parameter(8) - %constant.1144.clone.1 = f32[]{:T(128)} constant(0.9) - %broadcast.883.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%constant.1144.clone.1), dimensions={}, metadata={op_name="broadcast.81"} - %mul.1994.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%param_8.126, %broadcast.883.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.983.clone.1 = f32[2048,4]{0,1:T(4,128)S(1)} add(%mul.1995.clone.1, %mul.1994.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1322 = f32[]{:T(128)S(6)} parameter(2) - %div.920.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_2.1322), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %integer_pow.73.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%select_n.299.clone.1, %select_n.299.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} - %constant.1143.clone.1 = f32[]{:T(128)} constant(0.05) - %broadcast.882.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%constant.1143.clone.1), dimensions={}, metadata={op_name="broadcast.71"} - %mul.1993.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%integer_pow.73.clone.1, %broadcast.882.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.564 = f32[2048,4]{0,1:T(4,128)S(1)} parameter(4) - %constant.1142.clone.1 = f32[]{:T(128)} constant(0.95) - %broadcast.881.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%constant.1142.clone.1), dimensions={}, metadata={op_name="broadcast.70"} - %mul.1992.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%param_4.564, %broadcast.881.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.982.clone.1 = f32[2048,4]{0,1:T(4,128)S(1)} add(%mul.1993.clone.1, %mul.1992.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1564 = f32[]{:T(128)S(6)} parameter(1) - %div.919.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%param_1.1564), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.918.clone.1 = f32[2048,4]{0,1:T(4,128)} divide(%add.982.clone.1, %div.919.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %sqrt.70.clone.1 = f32[2048,4]{0,1:T(4,128)} sqrt(%div.918.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} - %constant.1141.clone.1 = f32[]{:T(128)} constant(1e-08) - %broadcast.879.clone.1 = f32[2048,4]{0,1:T(4,128)} broadcast(%constant.1141.clone.1), dimensions={}, metadata={op_name="broadcast.64"} - %add.981.clone.1 = f32[2048,4]{0,1:T(4,128)} add(%sqrt.70.clone.1, %broadcast.879.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %multiply.434.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%div.920.clone.1, %add.981.clone.1), metadata={op_name="multiply.53"} - %div.917.clone.1 = f32[2048,4]{0,1:T(4,128)} divide(%add.983.clone.1, %multiply.434.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1990.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%param_0.1376, %broadcast.884.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.980.clone.1 = f32[2048,4]{0,1:T(4,128)} add(%div.917.clone.1, %mul.1990.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1989.clone.1 = f32[2048,4]{0,1:T(4,128)} multiply(%mul.1991.clone.1, %add.980.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.979.clone.1 = f32[2048,4]{0,1:T(4,128)S(1)} add(%param_0.1376, %mul.1989.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.254 = f32[2048,4]{0,1:T(4,128)} multiply(%add.979.clone.1, %add.979.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1208 = f32[]{:T(128)} constant(0) - %reduce.208 = f32[]{:T(128)} reduce(%square.254, %constant.1208), dimensions={0,1}, to_apply=%region_63.68, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.210.clone.1 = f32[]{:T(128)} reduce(%integer_pow.73.clone.1, %constant.1208), dimensions={0,1}, to_apply=%region_48.53, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.153 = (f32[]{:T(128)}, f32[2048,4]{0,1:T(4,128)S(1)}, f32[2048,4]{0,1:T(4,128)S(1)}, f32[2048,4]{0,1:T(4,128)S(1)}, f32[]{:T(128)}) tuple(%reduce.208, %add.979.clone.1, %add.982.clone.1, %add.983.clone.1, %reduce.210.clone.1) -} - -%region_11.14 (reduce_sum.192: f32[], reduce_sum.193: f32[]) -> f32[] { - %reduce_sum.193 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.192 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.197 = f32[]{:T(128)} add(%reduce_sum.192, %reduce_sum.193), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.395 (param_0.1390: bf16[2048]) -> f32[] { - %param_0.1390 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(0) - %convert_element_type.1396 = f32[2048]{0:T(1024)} convert(%param_0.1390), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %square.257 = f32[2048]{0:T(1024)} multiply(%convert_element_type.1396, %convert_element_type.1396), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1222 = f32[]{:T(128)} constant(0) - ROOT %reduce.211 = f32[]{:T(128)} reduce(%square.257, %constant.1222), dimensions={0}, to_apply=%region_11.14, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} -} - -%region_59.64 (reduce_sum.401: f32[], reduce_sum.402: f32[]) -> f32[] { - %reduce_sum.402 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.401 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.403 = f32[]{:T(128)} add(%reduce_sum.401, %reduce_sum.402), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_44.49 (reduce_sum.323: f32[], reduce_sum.324: f32[]) -> f32[] { - %reduce_sum.324 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.323 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.325 = f32[]{:T(128)} add(%reduce_sum.323, %reduce_sum.324), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.396 (param_0.1380: f32[2048], param_1.1568: f32[], param_2.1326: f32[], param_3.930: f32[], param_4.568: f32[2048], param_5.480: f32[], param_6.370: bf16[2048], param_7.213: pred[], param_8.130: f32[2048]) -> (f32[], f32[2048], f32[2048], f32[2048], f32[]) { - %param_0.1380 = f32[2048]{0:T(1024)S(1)} parameter(0) - %param_3.930 = f32[]{:T(128)S(6)} parameter(3) - %mul.2022.clone.1 = f32[2048]{0:T(1024)} broadcast(%param_3.930), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.213 = pred[]{:T(512)S(6)} parameter(7) - %select_n.316.clone.1 = pred[2048]{0:T(1024)(128)(4,1)} broadcast(%param_7.213), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.370 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(6) - %convert_element_type.1411.clone.1 = f32[2048]{0:T(1024)} convert(%param_6.370), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %param_5.480 = f32[]{:T(128)} parameter(5) - %div.956.clone.1 = f32[2048]{0:T(1024)} broadcast(%param_5.480), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.955.clone.1 = f32[2048]{0:T(1024)} divide(%convert_element_type.1411.clone.1, %div.956.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.315.clone.1 = f32[2048]{0:T(1024)} select(%select_n.316.clone.1, %convert_element_type.1411.clone.1, %div.955.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %constant.1164.clone.1 = f32[]{:T(128)} constant(0.1) - %broadcast.900.clone.1 = f32[2048]{0:T(1024)} broadcast(%constant.1164.clone.1), dimensions={}, metadata={op_name="broadcast.86"} - %mul.2028.clone.1 = f32[2048]{0:T(1024)} multiply(%select_n.315.clone.1, %broadcast.900.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.130 = f32[2048]{0:T(1024)S(1)} parameter(8) - %constant.1168.clone.1 = f32[]{:T(128)} constant(0.9) - %mul.2029.clone.1 = f32[2048]{0:T(1024)} broadcast(%constant.1168.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.2027.clone.1 = f32[2048]{0:T(1024)} multiply(%param_8.130, %mul.2029.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.1005.clone.1 = f32[2048]{0:T(1024)S(1)} add(%mul.2028.clone.1, %mul.2027.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1326 = f32[]{:T(128)S(6)} parameter(2) - %div.952.clone.1 = f32[2048]{0:T(1024)} broadcast(%param_2.1326), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %integer_pow.77.clone.1 = f32[2048]{0:T(1024)} multiply(%select_n.315.clone.1, %select_n.315.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} - %constant.1167.clone.1 = f32[]{:T(128)} constant(0.05) - %mul.2026.clone.1 = f32[2048]{0:T(1024)} broadcast(%constant.1167.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.2024.clone.1 = f32[2048]{0:T(1024)} multiply(%integer_pow.77.clone.1, %mul.2026.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.568 = f32[2048]{0:T(1024)S(1)} parameter(4) - %constant.1166.clone.1 = f32[]{:T(128)} constant(0.95) - %mul.2025.clone.1 = f32[2048]{0:T(1024)} broadcast(%constant.1166.clone.1), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %mul.2023.clone.1 = f32[2048]{0:T(1024)} multiply(%param_4.568, %mul.2025.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.1004.clone.1 = f32[2048]{0:T(1024)S(1)} add(%mul.2024.clone.1, %mul.2023.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1568 = f32[]{:T(128)S(6)} parameter(1) - %div.951.clone.1 = f32[2048]{0:T(1024)} broadcast(%param_1.1568), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.950.clone.1 = f32[2048]{0:T(1024)} divide(%add.1004.clone.1, %div.951.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %sqrt.74.clone.1 = f32[2048]{0:T(1024)} sqrt(%div.950.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} - %constant.1165.clone.1 = f32[]{:T(128)} constant(1e-08) - %add.1003.clone.1 = f32[2048]{0:T(1024)} broadcast(%constant.1165.clone.1), dimensions={}, metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %add.1002.clone.1 = f32[2048]{0:T(1024)} add(%sqrt.74.clone.1, %add.1003.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %multiply.438.clone.1 = f32[2048]{0:T(1024)} multiply(%div.952.clone.1, %add.1002.clone.1), metadata={op_name="multiply.49"} - %div.949.clone.1 = f32[2048]{0:T(1024)} divide(%add.1005.clone.1, %multiply.438.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.2021.clone.1 = f32[2048]{0:T(1024)} multiply(%param_0.1380, %broadcast.900.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.1001.clone.1 = f32[2048]{0:T(1024)} add(%div.949.clone.1, %mul.2021.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.2020.clone.1 = f32[2048]{0:T(1024)} multiply(%mul.2022.clone.1, %add.1001.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.1000.clone.1 = f32[2048]{0:T(1024)S(1)} add(%param_0.1380, %mul.2020.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.258 = f32[2048]{0:T(1024)} multiply(%add.1000.clone.1, %add.1000.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1212 = f32[]{:T(128)} constant(0) - %reduce.212 = f32[]{:T(128)} reduce(%square.258, %constant.1212), dimensions={0}, to_apply=%region_59.64, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.213.clone.1 = f32[]{:T(128)} reduce(%integer_pow.77.clone.1, %constant.1212), dimensions={0}, to_apply=%region_44.49, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.156 = (f32[]{:T(128)}, f32[2048]{0:T(1024)S(1)}, f32[2048]{0:T(1024)S(1)}, f32[2048]{0:T(1024)S(1)}, f32[]{:T(128)}) tuple(%reduce.212, %add.1000.clone.1, %add.1004.clone.1, %add.1005.clone.1, %reduce.213.clone.1) -} - -%fused_computation.402 (param_0.1150: s32[512]) -> s32[1024] { - %constant.972 = s32[] constant(0), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} - %broadcast.815 = s32[1024]{0:T(1024)} broadcast(%constant.972), dimensions={}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} - %param_0.1150 = s32[512]{0:T(512)S(1)} parameter(0) - %constant.973 = s32[] constant(2147483647), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} - %pad.49 = s32[1024]{0:T(1024)} pad(%param_0.1150, %constant.973), padding=0_512, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} - %constant.971 = s32[] constant(151935), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} - %broadcast.814 = s32[1024]{0:T(1024)} broadcast(%constant.971), dimensions={}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} - ROOT %clamp.1 = s32[1024]{0:T(1024)} clamp(%broadcast.815, %pad.49, %broadcast.814), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/gather" stack_frame_id=0} -} - -%fused_computation.405 (param_0.1149: s32[4,128]) -> s32[512] { - %param_0.1149 = s32[4,128]{1,0:T(4,128)} parameter(0) - %constant.1065 = s32[]{:T(128)} constant(0) - %broadcast.834 = s32[4,128]{1,0:T(4,128)} broadcast(%constant.1065), dimensions={}, metadata={op_name="broadcast.95"} - %lt.32 = pred[4,128]{1,0:T(4,128)(4,1)} compare(%param_0.1149, %broadcast.834), direction=LT, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/lt" stack_frame_id=0} - %constant.1051 = s32[]{:T(128)} constant(151936) - %add.925 = s32[4,128]{1,0:T(4,128)} broadcast(%constant.1051), dimensions={}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/add" stack_frame_id=0} - %add.903 = s32[4,128]{1,0:T(4,128)} add(%param_0.1149, %add.925), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/add" stack_frame_id=0} - %select_n.178 = s32[4,128]{1,0:T(4,128)} select(%lt.32, %add.903, %param_0.1149), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/select_n" stack_frame_id=0} - ROOT %bitcast.409 = s32[512]{0:T(512)S(1)} bitcast(%select_n.178), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/broadcast_in_dim" stack_frame_id=0} -} - -%region_40.45 (reduce_sum.305: f32[], reduce_sum.309: f32[]) -> f32[] { - %reduce_sum.309 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.305 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.310 = f32[]{:T(128)} add(%reduce_sum.305, %reduce_sum.309), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_37.42 (reduce_sum.290: f32[], reduce_sum.291: f32[]) -> f32[] { - %reduce_sum.291 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.290 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.295 = f32[]{:T(128)} add(%reduce_sum.290, %reduce_sum.291), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.407 (param_0.1384: f32[4,128], param_1.1570: f32[4,128]) -> (f32[], f32[]) { - %param_0.1384 = f32[4,128]{1,0:T(4,128)} parameter(0) - %bitcast.413 = f32[128,4]{0,1:T(4,128)} bitcast(%param_0.1384), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %square.261 = f32[128,4]{0,1:T(4,128)} multiply(%bitcast.413, %bitcast.413), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1216 = f32[]{:T(128)} constant(0) - %reduce.214 = f32[]{:T(128)} reduce(%square.261, %constant.1216), dimensions={0,1}, to_apply=%region_40.45, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %param_1.1570 = f32[4,128]{1,0:T(4,128)} parameter(1) - %bitcast.417.clone.1 = f32[128,4]{0,1:T(4,128)} bitcast(%param_1.1570), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %square.264.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%bitcast.417.clone.1, %bitcast.417.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %reduce.215.clone.1 = f32[]{:T(128)} reduce(%square.264.clone.1, %constant.1216), dimensions={0,1}, to_apply=%region_37.42, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.170 = (f32[]{:T(128)}, f32[]{:T(128)}) tuple(%reduce.214, %reduce.215.clone.1) -} - -%region_72.77 (reduce_sum.470: f32[], reduce_sum.471: f32[]) -> f32[] { - %reduce_sum.471 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} - %reduce_sum.470 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} - ROOT %reduce_sum.472 = f32[]{:T(128)} add(%reduce_sum.470, %reduce_sum.471), metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +%fused_computation.424 (param_0.1209: bf16[1,128,1,64]) -> bf16[1,128,1,128] { + %param_0.1209 = bf16[1,128,1,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %constant.471.clone.2 = bf16[]{:T(256)} constant(-inf) + %pad.54 = bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1209, %constant.471.clone.2), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/concatenate" stack_frame_id=0} + %convert.201 = f32[1,128,1,128]{3,1,2,0:T(8,128)} convert(%pad.54) + %pad.53 = bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1209, %constant.471.clone.2), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/concatenate" stack_frame_id=0} + %convert.202 = f32[1,128,1,128]{3,1,2,0:T(8,128)} convert(%pad.53) + %maximum.46 = f32[1,128,1,128]{3,1,2,0:T(8,128)} maximum(%convert.201, %convert.202), metadata={op_name="jit(train_step)/concatenate" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","1","128"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + ROOT %convert.203 = bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)S(1)} convert(%maximum.46) +} + +%fused_computation.425 (param_0.1211: bf16[1,128,1,64]) -> bf16[1,128,1,128] { + %param_0.1211 = bf16[1,128,1,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %constant.471.clone.1 = bf16[]{:T(256)} constant(-inf) + %pad.56 = bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1211, %constant.471.clone.1), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/concatenate" stack_frame_id=0} + %convert.204 = f32[1,128,1,128]{3,1,2,0:T(8,128)} convert(%pad.56) + %pad.55 = bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1211, %constant.471.clone.1), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/concatenate" stack_frame_id=0} + %convert.205 = f32[1,128,1,128]{3,1,2,0:T(8,128)} convert(%pad.55) + %maximum.47 = f32[1,128,1,128]{3,1,2,0:T(8,128)} maximum(%convert.204, %convert.205), metadata={op_name="jit(train_step)/concatenate" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","1","128"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + ROOT %convert.206 = bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)S(1)} convert(%maximum.47) } -%region_58.63 (reduce_sum.395: f32[], reduce_sum.396: f32[]) -> f32[] { - %reduce_sum.396 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} - %reduce_sum.395 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} - ROOT %reduce_sum.400 = f32[]{:T(128)} add(%reduce_sum.395, %reduce_sum.396), metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +%fused_computation.493 (param_0.1516: f32[128,4]) -> bf16[4,128] { + %param_0.1516 = f32[128,4]{0,1:T(4,128)S(1)} parameter(0) + %bitcast.756 = f32[4,128]{1,0:T(4,128)} bitcast(%param_0.1516), metadata={op_name="jit(train_step)/jvp()/transpose" stack_frame_id=0} + ROOT %convert.176 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} convert(%bitcast.756) } -%fused_computation.410 (param_0.1391: bf16[4,128], param_1.1576: f32[4,128], param_2.1329: f32[4,128], param_3.932: s32[4,128]) -> (f32[], f32[], pred[4,128], f32[4,128]) { - %param_3.932 = s32[4,128]{1,0:T(4,128)S(1)} parameter(3) - %constant.1170.clone.1 = s32[]{:T(128)} constant(0) - %broadcast.901.clone.1 = s32[4,128]{1,0:T(4,128)} broadcast(%constant.1170.clone.1), dimensions={}, metadata={op_name="broadcast.95"} - %ne.6.clone.1 = pred[4,128]{1,0:T(4,128)(4,1)S(1)} compare(%param_3.932, %broadcast.901.clone.1), direction=NE, metadata={op_name="jit(train_step)/jvp()/ne" stack_frame_id=0} - %param_1.1576 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %log.16 = f32[4,128]{1,0:T(4,128)} log(%param_1.1576), metadata={op_name="jit(train_step)/jvp()/log" stack_frame_id=0} - %param_0.1391 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(0) - %reduce_max.15 = f32[4,128]{1,0:T(4,128)} convert(%param_0.1391), metadata={op_name="jit(train_step)/jvp()/reduce_max" stack_frame_id=0} - %add.927 = f32[4,128]{1,0:T(4,128)} add(%log.16, %reduce_max.15), metadata={op_name="jit(train_step)/jvp()/add" stack_frame_id=0} - %square.269 = f32[4,128]{1,0:T(4,128)} multiply(%add.927, %add.927), metadata={op_name="jit(train_step)/jvp()/square" stack_frame_id=0} - %constant.1224 = f32[]{:T(128)} constant(0) - %broadcast.831 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1224), dimensions={}, metadata={op_name="broadcast.99"} - %mul.1913 = f32[4,128]{1,0:T(4,128)} multiply(%square.269, %broadcast.831), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} - %mul.1893 = f32[4,128]{1,0:T(4,128)} select(%ne.6.clone.1, %mul.1913, %broadcast.831), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} - %reduce.216 = f32[]{:T(128)} reduce(%mul.1893, %constant.1224), dimensions={0,1}, to_apply=%region_72.77, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} - %param_2.1329 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %neg.115.clone.1 = f32[4,128]{1,0:T(4,128)} negate(%param_2.1329), metadata={op_name="jit(train_step)/jvp()/neg" stack_frame_id=0} - %add.904.clone.1 = f32[4,128]{1,0:T(4,128)} add(%neg.115.clone.1, %mul.1913), metadata={op_name="jit(train_step)/jvp()/add" stack_frame_id=0} - %mul.1894.clone.1 = f32[4,128]{1,0:T(4,128)} select(%ne.6.clone.1, %add.904.clone.1, %broadcast.831), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} - %reduce.219.clone.1 = f32[]{:T(128)} reduce(%mul.1894.clone.1, %constant.1224), dimensions={0,1}, to_apply=%region_58.63, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} - %mul.1911.clone.1 = f32[4,128]{1,0:T(4,128)} multiply(%add.927, %broadcast.831), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} - %constant.1068.clone.1 = f32[]{:T(128)} constant(1) - %add.922.clone.1 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1068.clone.1), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/add" stack_frame_id=0} - %add.915.clone.1 = f32[4,128]{1,0:T(4,128)S(1)} add(%mul.1911.clone.1, %add.922.clone.1), metadata={op_name="jit(train_step)/transpose(jvp())/add" stack_frame_id=0} - ROOT %tuple.157 = (f32[]{:T(128)}, f32[]{:T(128)}, pred[4,128]{1,0:T(4,128)(4,1)S(1)}, f32[4,128]{1,0:T(4,128)S(1)}) tuple(%reduce.216, %reduce.219.clone.1, %ne.6.clone.1, %add.915.clone.1) -} - -%region_69.74 (reduce_sum.452: f32[], reduce_sum.456: f32[]) -> f32[] { - %reduce_sum.456 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.452 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.457 = f32[]{:T(128)} add(%reduce_sum.452, %reduce_sum.456), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_54.59 (reduce_sum.374: f32[], reduce_sum.375: f32[]) -> f32[] { - %reduce_sum.375 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.374 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.379 = f32[]{:T(128)} add(%reduce_sum.374, %reduce_sum.375), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.411 (param_0.1370: f32[128,4], param_1.1558: f32[], param_2.1316: f32[], param_3.920: f32[], param_4.558: f32[128,4], param_5.470: f32[], param_6.360: f32[4,128], param_7.203: pred[], param_8.120: f32[128,4]) -> (f32[], f32[128,4], f32[128,4], f32[128,4], f32[]) { - %param_0.1370 = f32[128,4]{0,1:T(4,128)S(1)} parameter(0) - %param_3.920 = f32[]{:T(128)S(6)} parameter(3) - %mul.1943.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_3.920), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.203 = pred[]{:T(512)S(6)} parameter(7) - %select_n.276.clone.1 = pred[128,4]{0,1:T(4,128)(4,1)} broadcast(%param_7.203), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.360 = f32[4,128]{1,0:T(4,128)} parameter(6) - %bitcast.468.clone.1 = f32[128,4]{0,1:T(4,128)} bitcast(%param_6.360), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.470 = f32[]{:T(128)} parameter(5) - %div.876.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_5.470), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.875.clone.1 = f32[128,4]{0,1:T(4,128)} divide(%bitcast.468.clone.1, %div.876.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.275.clone.1 = f32[128,4]{0,1:T(4,128)} select(%select_n.276.clone.1, %bitcast.468.clone.1, %div.875.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %constant.1104.clone.1 = f32[]{:T(128)} constant(0.1) - %broadcast.856.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%constant.1104.clone.1), dimensions={}, metadata={op_name="broadcast.78"} - %mul.1947.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%select_n.275.clone.1, %broadcast.856.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.120 = f32[128,4]{0,1:T(4,128)S(1)} parameter(8) - %constant.1108.clone.1 = f32[]{:T(128)} constant(0.9) - %broadcast.855.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%constant.1108.clone.1), dimensions={}, metadata={op_name="broadcast.77"} - %mul.1946.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%param_8.120, %broadcast.855.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.951.clone.1 = f32[128,4]{0,1:T(4,128)S(1)} add(%mul.1947.clone.1, %mul.1946.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1316 = f32[]{:T(128)S(6)} parameter(2) - %div.872.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_2.1316), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %integer_pow.67.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%select_n.275.clone.1, %select_n.275.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} - %constant.1107.clone.1 = f32[]{:T(128)} constant(0.05) - %broadcast.854.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%constant.1107.clone.1), dimensions={}, metadata={op_name="broadcast.67"} - %mul.1945.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%integer_pow.67.clone.1, %broadcast.854.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.558 = f32[128,4]{0,1:T(4,128)S(1)} parameter(4) - %constant.1106.clone.1 = f32[]{:T(128)} constant(0.95) - %broadcast.853.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%constant.1106.clone.1), dimensions={}, metadata={op_name="broadcast.66"} - %mul.1944.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%param_4.558, %broadcast.853.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.950.clone.1 = f32[128,4]{0,1:T(4,128)S(1)} add(%mul.1945.clone.1, %mul.1944.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1558 = f32[]{:T(128)S(6)} parameter(1) - %div.871.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_1.1558), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.870.clone.1 = f32[128,4]{0,1:T(4,128)} divide(%add.950.clone.1, %div.871.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %sqrt.64.clone.1 = f32[128,4]{0,1:T(4,128)} sqrt(%div.870.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} - %constant.1105.clone.1 = f32[]{:T(128)} constant(1e-08) - %broadcast.851.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%constant.1105.clone.1), dimensions={}, metadata={op_name="broadcast.62"} - %add.949.clone.1 = f32[128,4]{0,1:T(4,128)} add(%sqrt.64.clone.1, %broadcast.851.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %multiply.428.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%div.872.clone.1, %add.949.clone.1), metadata={op_name="multiply.59"} - %div.869.clone.1 = f32[128,4]{0,1:T(4,128)} divide(%add.951.clone.1, %multiply.428.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1942.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%param_0.1370, %broadcast.856.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.948.clone.1 = f32[128,4]{0,1:T(4,128)} add(%div.869.clone.1, %mul.1942.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1941.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%mul.1943.clone.1, %add.948.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.947.clone.1 = f32[128,4]{0,1:T(4,128)S(1)} add(%param_0.1370, %mul.1941.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.265 = f32[128,4]{0,1:T(4,128)} multiply(%add.947.clone.1, %add.947.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1202 = f32[]{:T(128)} constant(0) - %reduce.217 = f32[]{:T(128)} reduce(%square.265, %constant.1202), dimensions={0,1}, to_apply=%region_69.74, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.221.clone.1 = f32[]{:T(128)} reduce(%integer_pow.67.clone.1, %constant.1202), dimensions={0,1}, to_apply=%region_54.59, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.159 = (f32[]{:T(128)}, f32[128,4]{0,1:T(4,128)S(1)}, f32[128,4]{0,1:T(4,128)S(1)}, f32[128,4]{0,1:T(4,128)S(1)}, f32[]{:T(128)}) tuple(%reduce.217, %add.947.clone.1, %add.950.clone.1, %add.951.clone.1, %reduce.221.clone.1) -} - -%region_66.71 (reduce_sum.437: f32[], reduce_sum.438: f32[]) -> f32[] { - %reduce_sum.438 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.437 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.442 = f32[]{:T(128)} add(%reduce_sum.437, %reduce_sum.438), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%region_51.56 (reduce_sum.359: f32[], reduce_sum.360: f32[]) -> f32[] { - %reduce_sum.360 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/reduce_sum"} - %reduce_sum.359 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/reduce_sum"} - ROOT %reduce_sum.361 = f32[]{:T(128)} add(%reduce_sum.359, %reduce_sum.360), metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.412 (param_0.1373: f32[128,4], param_1.1561: f32[], param_2.1319: f32[], param_3.923: f32[], param_4.561: f32[128,4], param_5.473: f32[], param_6.363: f32[4,128], param_7.206: pred[], param_8.123: f32[128,4]) -> (f32[], f32[128,4], f32[128,4], f32[128,4], f32[]) { - %param_0.1373 = f32[128,4]{0,1:T(4,128)S(1)} parameter(0) - %param_3.923 = f32[]{:T(128)S(6)} parameter(3) - %mul.1970.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_3.923), dimensions={}, metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_7.206 = pred[]{:T(512)S(6)} parameter(7) - %select_n.288.clone.1 = pred[128,4]{0,1:T(4,128)(4,1)} broadcast(%param_7.206), dimensions={}, metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %param_6.363 = f32[4,128]{1,0:T(4,128)} parameter(6) - %bitcast.474.clone.1 = f32[128,4]{0,1:T(4,128)} bitcast(%param_6.363), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - %param_5.473 = f32[]{:T(128)} parameter(5) - %div.900.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_5.473), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.899.clone.1 = f32[128,4]{0,1:T(4,128)} divide(%bitcast.474.clone.1, %div.900.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %select_n.287.clone.1 = f32[128,4]{0,1:T(4,128)} select(%select_n.288.clone.1, %bitcast.474.clone.1, %div.899.clone.1), metadata={op_name="jit(train_step)/select_n" stack_frame_id=0} - %constant.1122.clone.1 = f32[]{:T(128)} constant(0.1) - %broadcast.866.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%constant.1122.clone.1), dimensions={}, metadata={op_name="broadcast.78"} - %mul.1974.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%select_n.287.clone.1, %broadcast.866.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_8.123 = f32[128,4]{0,1:T(4,128)S(1)} parameter(8) - %constant.1126.clone.1 = f32[]{:T(128)} constant(0.9) - %broadcast.865.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%constant.1126.clone.1), dimensions={}, metadata={op_name="broadcast.77"} - %mul.1973.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%param_8.123, %broadcast.865.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.968.clone.1 = f32[128,4]{0,1:T(4,128)S(1)} add(%mul.1974.clone.1, %mul.1973.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_2.1319 = f32[]{:T(128)S(6)} parameter(2) - %div.896.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_2.1319), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %integer_pow.70.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%select_n.287.clone.1, %select_n.287.clone.1), metadata={op_name="jit(train_step)/integer_pow" stack_frame_id=0} - %constant.1125.clone.1 = f32[]{:T(128)} constant(0.05) - %broadcast.864.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%constant.1125.clone.1), dimensions={}, metadata={op_name="broadcast.67"} - %mul.1972.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%integer_pow.70.clone.1, %broadcast.864.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %param_4.561 = f32[128,4]{0,1:T(4,128)S(1)} parameter(4) - %constant.1124.clone.1 = f32[]{:T(128)} constant(0.95) - %broadcast.863.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%constant.1124.clone.1), dimensions={}, metadata={op_name="broadcast.66"} - %mul.1971.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%param_4.561, %broadcast.863.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.967.clone.1 = f32[128,4]{0,1:T(4,128)S(1)} add(%mul.1972.clone.1, %mul.1971.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %param_1.1561 = f32[]{:T(128)S(6)} parameter(1) - %div.895.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%param_1.1561), dimensions={}, metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %div.894.clone.1 = f32[128,4]{0,1:T(4,128)} divide(%add.967.clone.1, %div.895.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %sqrt.67.clone.1 = f32[128,4]{0,1:T(4,128)} sqrt(%div.894.clone.1), metadata={op_name="jit(train_step)/sqrt" stack_frame_id=0} - %constant.1123.clone.1 = f32[]{:T(128)} constant(1e-08) - %broadcast.861.clone.1 = f32[128,4]{0,1:T(4,128)} broadcast(%constant.1123.clone.1), dimensions={}, metadata={op_name="broadcast.62"} - %add.966.clone.1 = f32[128,4]{0,1:T(4,128)} add(%sqrt.67.clone.1, %broadcast.861.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %multiply.431.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%div.896.clone.1, %add.966.clone.1), metadata={op_name="multiply.56"} - %div.893.clone.1 = f32[128,4]{0,1:T(4,128)} divide(%add.968.clone.1, %multiply.431.clone.1), metadata={op_name="jit(train_step)/div" stack_frame_id=0} - %mul.1969.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%param_0.1373, %broadcast.866.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.965.clone.1 = f32[128,4]{0,1:T(4,128)} add(%div.893.clone.1, %mul.1969.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %mul.1968.clone.1 = f32[128,4]{0,1:T(4,128)} multiply(%mul.1970.clone.1, %add.965.clone.1), metadata={op_name="jit(train_step)/mul" stack_frame_id=0} - %add.964.clone.1 = f32[128,4]{0,1:T(4,128)S(1)} add(%param_0.1373, %mul.1968.clone.1), metadata={op_name="jit(train_step)/add" stack_frame_id=0} - %square.266 = f32[128,4]{0,1:T(4,128)} multiply(%add.964.clone.1, %add.964.clone.1), metadata={op_name="jit(train_step)/square" stack_frame_id=0} - %constant.1205 = f32[]{:T(128)} constant(0) - %reduce.218 = f32[]{:T(128)} reduce(%square.266, %constant.1205), dimensions={0,1}, to_apply=%region_66.71, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - %reduce.222.clone.1 = f32[]{:T(128)} reduce(%integer_pow.70.clone.1, %constant.1205), dimensions={0,1}, to_apply=%region_51.56, metadata={op_name="jit(train_step)/reduce_sum" stack_frame_id=0} - ROOT %tuple.160 = (f32[]{:T(128)}, f32[128,4]{0,1:T(4,128)S(1)}, f32[128,4]{0,1:T(4,128)S(1)}, f32[128,4]{0,1:T(4,128)S(1)}, f32[]{:T(128)}) tuple(%reduce.218, %add.964.clone.1, %add.967.clone.1, %add.968.clone.1, %reduce.222.clone.1) -} - -%fused_computation.421 (param_0.1201: f32[4,128], param_1.1323: f32[4,128]) -> f32[4,128] { - %param_0.1201 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %param_1.1323 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %constant.1045 = f32[]{:T(128)} constant(0.00048828125) - %broadcast.837 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1045), dimensions={}, metadata={op_name="broadcast.399"} - %div.767 = f32[4,128]{1,0:T(4,128)} multiply(%param_1.1323, %broadcast.837), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/div" stack_frame_id=0} - %constant.1043 = f32[]{:T(128)} constant(1e-06) - %add.935 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1043), dimensions={}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/add" stack_frame_id=0} - %add.934 = f32[4,128]{1,0:T(4,128)} add(%div.767, %add.935), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/add" stack_frame_id=0} - %rsqrt.168 = f32[4,128]{1,0:T(4,128)} rsqrt(%add.934), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/rsqrt" stack_frame_id=0} - %div.754 = f32[4,128]{1,0:T(4,128)} divide(%rsqrt.168, %add.934), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/div" stack_frame_id=0} - %constant.1040 = f32[]{:T(128)} constant(-0.5) - %mul.1919 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1040), dimensions={}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %mul.1910 = f32[4,128]{1,0:T(4,128)} multiply(%div.754, %mul.1919), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %mul.1909 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.1201, %mul.1910), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %constant.1039 = f32[]{:T(128)} constant(0.0009765625) - %mul.1918 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1039), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - ROOT %mul.1908 = f32[4,128]{1,0:T(4,128)S(1)} multiply(%mul.1909, %mul.1918), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} +%fused_computation.387 (param_0.1098: f32[512,4,8,128]) -> bf16[4,512,8,128] { + %param_0.1098 = f32[512,4,8,128]{3,2,1,0:T(8,128)S(1)} parameter(0) + %copy.213 = bf16[512,4,8,128]{3,0,2,1:T(8,128)(2,1)} copy(%param_0.1098), sharding={devices=[4,1,1,1]<=[4]}, metadata={op_name="state[\'model\'][\'decoder\'][\'layers\'][\'self_attention\'][\'value\'][\'kernel\'].value"} + ROOT %bitcast.568 = bf16[4,512,8,128]{3,1,2,0:T(8,128)(2,1)} bitcast(%copy.213), metadata={op_name="jit(train_step)/jvp()/transpose" stack_frame_id=0} +} + +%fused_computation.370 (param_0.1066: f32[16,4,128,512]) -> bf16[4,16,128,512] { + %param_0.1066 = f32[16,4,128,512]{3,2,1,0:T(8,128)S(1)} parameter(0) + %copy.212 = bf16[16,4,128,512]{3,2,0,1:T(8,128)(2,1)} copy(%param_0.1066), sharding={devices=[1,1,1,4]<=[4]}, metadata={op_name="state[\'model\'][\'decoder\'][\'layers\'][\'self_attention\'][\'out\'][\'kernel\'].value"} + ROOT %bitcast.559 = bf16[4,16,128,512]{3,2,1,0:T(8,128)(2,1)} bitcast(%copy.212), metadata={op_name="jit(train_step)/jvp()/transpose" stack_frame_id=0} +} + +%fused_computation.492 (param_0.1514: f32[2048,4]) -> bf16[4,2048] { + %param_0.1514 = f32[2048,4]{0,1:T(4,128)S(1)} parameter(0) + %bitcast.754 = f32[4,2048]{1,0:T(4,128)} bitcast(%param_0.1514), metadata={op_name="jit(train_step)/jvp()/transpose" stack_frame_id=0} + ROOT %convert.174 = bf16[4,2048]{1,0:T(4,128)(2,1)S(1)} convert(%bitcast.754) +} + +%fused_computation.347 (param_0.1017: f32[512,4,6144]) -> bf16[4,512,6144] { + %param_0.1017 = f32[512,4,6144]{2,1,0:T(4,128)} parameter(0) + %copy.211 = bf16[512,4,6144]{2,0,1:T(8,128)(2,1)} copy(%param_0.1017), sharding={devices=[4,1,1]<=[4]}, metadata={op_name="state[\'model\'][\'decoder\'][\'layers\'][\'mlp\'][\'wi_0\'][\'kernel\'].value"} + ROOT %bitcast.547 = bf16[4,512,6144]{2,1,0:T(8,128)(2,1)} bitcast(%copy.211), metadata={op_name="jit(train_step)/jvp()/transpose" stack_frame_id=0} +} + +%fused_computation.346 (param_0.1015: f32[512,4,6144]) -> bf16[4,512,6144] { + %param_0.1015 = f32[512,4,6144]{2,1,0:T(4,128)} parameter(0) + %copy.210 = bf16[512,4,6144]{2,0,1:T(8,128)(2,1)} copy(%param_0.1015), sharding={devices=[4,1,1]<=[4]}, metadata={op_name="state[\'model\'][\'decoder\'][\'layers\'][\'mlp\'][\'wi_1\'][\'kernel\'].value"} + ROOT %bitcast.546 = bf16[4,512,6144]{2,1,0:T(8,128)(2,1)} bitcast(%copy.210), metadata={op_name="jit(train_step)/jvp()/transpose" stack_frame_id=0} } -%region_0.1 (reduce_sum.137: s32[], reduce_sum.138: s32[]) -> s32[] { - %reduce_sum.138 = s32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} - %reduce_sum.137 = s32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} - ROOT %reduce_sum.139 = s32[]{:T(128)} add(%reduce_sum.137, %reduce_sum.138), metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[{"indices":["0","2"]}]}} +%fused_computation.345 (param_0.1013: f32[6144,4,512]) -> bf16[4,6144,512] { + %param_0.1013 = f32[6144,4,512]{2,1,0:T(4,128)S(1)} parameter(0) + %copy.209 = bf16[6144,4,512]{2,0,1:T(8,128)(2,1)} copy(%param_0.1013), sharding={devices=[1,1,4]<=[4]}, metadata={op_name="state[\'model\'][\'decoder\'][\'layers\'][\'mlp\'][\'wo\'][\'kernel\'].value"} + ROOT %bitcast.545 = bf16[4,6144,512]{2,1,0:T(8,128)(2,1)} bitcast(%copy.209), metadata={op_name="jit(train_step)/jvp()/transpose" stack_frame_id=0} +} + +%fused_computation.60.clone.1 (param_0.1806: bf16[4,1,128,2048], param_1.1890: s32[], param_2.1442: bf16[1,128,2048]) -> bf16[4,1,128,2048] { + %param_0.1806 = bf16[4,1,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(0) + %param_2.1442 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %bitcast.1011 = bf16[1,1,128,2048]{3,2,1,0:T(8,128)(2,1)} bitcast(%param_2.1442), metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_update_index_in_dim)/broadcast_in_dim" stack_frame_id=0} + %param_1.1890 = s32[]{:T(128)S(6)} parameter(1) + %constant.420.clone.35 = s32[]{:T(128)} constant(0) + ROOT %dynamic-update-slice.32 = bf16[4,1,128,2048]{3,2,1,0:T(8,128)(2,1)} dynamic-update-slice(%param_0.1806, %bitcast.1011, %param_1.1890, %constant.420.clone.35, %constant.420.clone.35, /*index=5*/%constant.420.clone.35), metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_update_index_in_dim)/dynamic_update_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} +} + +%fused_computation.22.clone.1 (param_0.1831: bf16[4,16,128,512], param_1.1906: s32[]) -> bf16[1,16,128,512] { + %param_0.1831 = bf16[4,16,128,512]{3,2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1906 = s32[]{:T(128)S(6)} parameter(1) + %constant.420.clone.41 = s32[]{:T(128)} constant(0) + ROOT %dynamic-slice.210 = bf16[1,16,128,512]{3,2,1,0:T(8,128)(2,1)S(1)} dynamic-slice(%param_0.1831, %param_1.1906, %constant.420.clone.41, %constant.420.clone.41, %constant.420.clone.41), dynamic_slice_sizes={1,16,128,512}, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_index_in_dim)/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} +} + +%fused_computation.23.clone.1 (param_0.1823: bf16[4,512,16,128], param_1.1901: s32[]) -> bf16[1,512,16,128] { + %param_0.1823 = bf16[4,512,16,128]{1,3,2,0:T(8,128)(2,1)S(1)} parameter(0) + %param_1.1901 = s32[]{:T(128)S(6)} parameter(1) + %constant.420.clone.40 = s32[]{:T(128)} constant(0) + ROOT %dynamic-slice.209 = bf16[1,512,16,128]{1,3,2,0:T(8,128)(2,1)S(1)} dynamic-slice(%param_0.1823, %param_1.1901, %constant.420.clone.40, %constant.420.clone.40, %constant.420.clone.40), dynamic_slice_sizes={1,512,16,128}, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_index_in_dim)/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} +} + +%region_2.3 (reduce_sum.143: f32[], reduce_sum.144: f32[]) -> f32[] { + %reduce_sum.143 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/reduce_sum"} + %reduce_sum.144 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/reduce_sum"} + ROOT %reduce_sum.145 = f32[]{:T(128)} add(%reduce_sum.143, %reduce_sum.144), metadata={op_name="reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%fused_computation.38.clone.1 (param_0.1808: bf16[1,128,2048]) -> f32[128] { + %param_0.1808 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.1516 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%param_0.1808), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %square.281 = f32[1,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1516, %convert_element_type.1516), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/square" stack_frame_id=0} + %constant.421.clone.13 = f32[]{:T(128)} constant(0) + ROOT %reduce.276 = f32[128]{0:T(128)S(1)} reduce(%square.281, %constant.421.clone.13), dimensions={0,2}, to_apply=%region_2.3, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/reduce_sum" stack_frame_id=0} +} + +%fused_computation.83.clone.1 (param_0.1809: f32[128]) -> f32[128] { + %param_0.1809 = f32[128]{0:T(128)S(1)} parameter(0) + %constant.422.clone.9 = f32[]{:T(128)} constant(0.00048828125) + %broadcast.1093 = f32[128]{0:T(128)} broadcast(%constant.422.clone.9), dimensions={}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call" stack_frame_id=0} + %div.1030 = f32[128]{0:T(128)} multiply(%param_0.1809, %broadcast.1093), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/div" stack_frame_id=0} + %constant.423.clone.13 = f32[]{:T(128)} constant(1e-06) + %broadcast.1092 = f32[128]{0:T(128)} broadcast(%constant.423.clone.13), dimensions={}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call" stack_frame_id=0} + %add.1061 = f32[128]{0:T(128)} add(%div.1030, %broadcast.1092), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/add" stack_frame_id=0} + ROOT %rsqrt.222 = f32[128]{0:T(128)S(1)} rsqrt(%add.1061), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/rsqrt" stack_frame_id=0} +} + +%convert_element_type.932.reduce_sub_computation (lhs.3: bf16[], rhs.3: bf16[]) -> bf16[] { + %lhs.3 = bf16[] parameter(0) + %rhs.3 = bf16[] parameter(1) + ROOT %add.795 = bf16[] add(%lhs.3, %rhs.3), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%convert_element_type.920.reduce_sub_computation (lhs.2: bf16[], rhs.2: bf16[]) -> bf16[] { + %lhs.2 = bf16[] parameter(0) + %rhs.2 = bf16[] parameter(1) + ROOT %add.794 = bf16[] add(%lhs.2, %rhs.2), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%fused_computation.73.clone.1 (param_0.1807: bf16[4,2048], param_1.1891: s32[], param_2.1443: bf16[4,2048]) -> (bf16[2048], bf16[2048]) { + %param_0.1807 = bf16[4,2048]{1,0:T(4,128)(2,1)S(1)} parameter(0) + %param_1.1891 = s32[]{:T(128)S(6)} parameter(1) + %constant.420.clone.36 = s32[]{:T(128)} constant(0) + %dynamic_slice.301 = bf16[1,2048]{1,0:T(2,128)(2,1)} dynamic-slice(%param_0.1807, %param_1.1891, %constant.420.clone.36), dynamic_slice_sizes={1,2048}, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_index_in_dim)/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %constant.1134 = bf16[]{:T(256)} constant(-0), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %reduce.275 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} reduce(%dynamic_slice.301, %constant.1134), dimensions={0}, to_apply=%convert_element_type.932.reduce_sub_computation, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %param_2.1443 = bf16[4,2048]{1,0:T(4,128)(2,1)S(1)} parameter(2) + %dynamic_slice.272.clone.3 = bf16[1,2048]{1,0:T(2,128)(2,1)} dynamic-slice(%param_2.1443, %param_1.1891, %constant.420.clone.36), dynamic_slice_sizes={1,2048}, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_index_in_dim)/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %reduce.181.clone.3 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} reduce(%dynamic_slice.272.clone.3, %constant.1134), dimensions={0}, to_apply=%convert_element_type.920.reduce_sub_computation, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + ROOT %tuple.240 = (bf16[2048]{0:T(1024)(128)(2,1)S(1)}, bf16[2048]{0:T(1024)(128)(2,1)S(1)}) tuple(%reduce.275, %reduce.181.clone.3) +} + +%fused_computation.53.clone.2.clone.clone.clone.1 (param_0.1825: bf16[1,128,2048], param_1.1902: f32[128], param_2.1451: bf16[2048]) -> bf16[128,2048,1] { + %param_2.1451 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(2) + %dot_general.636 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1451), dimensions={2}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/...k,k->...k/dot_general" stack_frame_id=0} + %convert.374 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%dot_general.636) + %param_0.1825 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.1524 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%param_0.1825), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %param_1.1902 = f32[128]{0:T(128)S(1)} parameter(1) + %mul.2392 = f32[1,128,2048]{2,1,0:T(8,128)} broadcast(%param_1.1902), dimensions={1}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %mul.2391 = f32[1,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1524, %mul.2392), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %convert_element_type.1523 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.2391), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %convert.375 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%convert_element_type.1523) + %dot_general.635 = f32[1,128,2048]{2,1,0:T(8,128)} multiply(%convert.374, %convert.375), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/...k,k->...k/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","2048"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %convert.376 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} convert(%dot_general.635) + ROOT %bitcast.1020 = bf16[128,2048,1]{1,0,2:T(8,128)(2,1)} bitcast(%convert.376) +} + +%fused_computation.17.clone.clone.clone.1 (param_0.1824: bf16[1,2048,16,128]) -> bf16[2048,16,128] { + %param_0.1824 = bf16[1,2048,16,128]{1,3,2,0:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.1019 = bf16[2048,16,128]{0,2,1:T(8,128)(2,1)} bitcast(%param_0.1824), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} +} + +%region_3.4 (reduce_sum.146: f32[], reduce_sum.150: f32[]) -> f32[] { + %reduce_sum.146 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/reduce_sum"} + %reduce_sum.150 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/reduce_sum"} + ROOT %reduce_sum.151 = f32[]{:T(128)} add(%reduce_sum.146, %reduce_sum.150), metadata={op_name="reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%fused_computation.44.clone.1 (param_0.1826: bf16[1,2048,16,128], param_1.1903: bf16[1,128,2048], param_2.1452: f32[128], param_3.954: bf16[2048]) -> (f32[128,16], bf16[128,16,128]) { + %param_1.1903 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %param_2.1452 = f32[128]{0:T(128)S(1)} parameter(2) + %param_3.954 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(3) + %fusion.91.clone.3 = bf16[128,2048,1]{1,0,2:T(8,128)(2,1)} fusion(%param_1.1903, %param_2.1452, %param_3.954), kind=kLoop, calls=%fused_computation.53.clone.2.clone.clone.clone.1 + %param_0.1826 = bf16[1,2048,16,128]{1,3,2,0:T(8,128)(2,1)S(1)} parameter(0) + %fusion.22.clone.3 = bf16[2048,16,128]{0,2,1:T(8,128)(2,1)} fusion(%param_0.1826), kind=kLoop, calls=%fused_computation.17.clone.clone.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} + %convolution.34.clone.3 = bf16[128,16,128]{2,0,1:T(8,128)(2,1)S(1)} convolution(%fusion.91.clone.3, %fusion.22.clone.3), window={size=16 pad=15_15 rhs_reversal=1}, dim_labels=bf0_i0o->b0f, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} + %convert_element_type.1525 = f32[128,16,128]{2,0,1:T(8,128)} convert(%convolution.34.clone.3), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %bitcast.1021 = f32[1,128,16,128]{3,1,2,0:T(8,128)} bitcast(%convert_element_type.1525), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %square.283 = f32[1,128,16,128]{3,1,2,0:T(8,128)} multiply(%bitcast.1021, %bitcast.1021), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/square" stack_frame_id=0} + %constant.421.clone.15 = f32[]{:T(128)} constant(0) + %reduce.278 = f32[128,16]{0,1:T(8,128)S(1)} reduce(%square.283, %constant.421.clone.15), dimensions={0,3}, to_apply=%region_3.4, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/reduce_sum" stack_frame_id=0} + ROOT %tuple.244 = (f32[128,16]{0,1:T(8,128)S(1)}, bf16[128,16,128]{2,0,1:T(8,128)(2,1)S(1)}) tuple(%reduce.278, %convolution.34.clone.3) +} + +%fused_computation.72.clone.1 (param_0.1827: f32[128,16]) -> f32[128,16] { + %param_0.1827 = f32[128,16]{0,1:T(8,128)S(1)} parameter(0) + %constant.424.clone.10 = f32[]{:T(128)} constant(0.0078125) + %broadcast.1097 = f32[128,16]{0,1:T(8,128)} broadcast(%constant.424.clone.10), dimensions={}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/div" stack_frame_id=0} + %div.1032 = f32[128,16]{0,1:T(8,128)} multiply(%param_0.1827, %broadcast.1097), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/div" stack_frame_id=0} + %constant.423.clone.15 = f32[]{:T(128)} constant(1e-06) + %broadcast.1096 = f32[128,16]{0,1:T(8,128)} broadcast(%constant.423.clone.15), dimensions={}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/add" stack_frame_id=0} + %add.1064 = f32[128,16]{0,1:T(8,128)} add(%div.1032, %broadcast.1096), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/add" stack_frame_id=0} + ROOT %rsqrt.224 = f32[128,16]{0,1:T(8,128)S(1)} rsqrt(%add.1064), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/rsqrt" stack_frame_id=0} +} + +%fused_computation.79.clone.1 (param_0.1814: bf16[4,128], param_1.1895: s32[], param_2.1446: bf16[4,128]) -> (bf16[128], bf16[128]) { + %param_0.1814 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(0) + %param_1.1895 = s32[]{:T(128)S(6)} parameter(1) + %constant.420.clone.38 = s32[]{:T(128)} constant(0) + %dynamic_slice.302 = bf16[1,128]{1,0:T(2,128)(2,1)} dynamic-slice(%param_0.1814, %param_1.1895, %constant.420.clone.38), dynamic_slice_sizes={1,128}, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_index_in_dim)/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %bitcast.1015 = bf16[128]{0:T(256)(128)(2,1)S(1)} bitcast(%dynamic_slice.302), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %param_2.1446 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} parameter(2) + %dynamic_slice.282.clone.3 = bf16[1,128]{1,0:T(2,128)(2,1)} dynamic-slice(%param_2.1446, %param_1.1895, %constant.420.clone.38), dynamic_slice_sizes={1,128}, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_index_in_dim)/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %bitcast.278.clone.3 = bf16[128]{0:T(256)(128)(2,1)S(1)} bitcast(%dynamic_slice.282.clone.3), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + ROOT %tuple.241 = (bf16[128]{0:T(256)(128)(2,1)S(1)}, bf16[128]{0:T(256)(128)(2,1)S(1)}) tuple(%bitcast.1015, %bitcast.278.clone.3) +} + +%fused_computation.49.clone.2 (param_0.1828: bf16[128,16,128], param_1.1904: f32[128,16], param_2.1453: bf16[128]) -> bf16[1,128,16,128] { + %param_2.1453 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(2) + %dot_general.638 = bf16[1,128,16,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_2.1453), dimensions={3}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/...k,k->...k/dot_general" stack_frame_id=0} + %convert.387 = f32[1,128,16,128]{3,1,2,0:T(8,128)} convert(%dot_general.638) + %param_0.1828 = bf16[128,16,128]{2,0,1:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.1527 = f32[128,16,128]{2,0,1:T(8,128)} convert(%param_0.1828), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %bitcast.1022 = f32[1,128,16,128]{3,1,2,0:T(8,128)} bitcast(%convert_element_type.1527), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %param_1.1904 = f32[128,16]{0,1:T(8,128)S(1)} parameter(1) + %mul.2394 = f32[1,128,16,128]{3,1,2,0:T(8,128)} broadcast(%param_1.1904), dimensions={1,2}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %mul.2393 = f32[1,128,16,128]{3,1,2,0:T(8,128)} multiply(%bitcast.1022, %mul.2394), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %convert_element_type.1526 = bf16[1,128,16,128]{3,1,2,0:T(8,128)(2,1)} convert(%mul.2393), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %convert.388 = f32[1,128,16,128]{3,1,2,0:T(8,128)} convert(%convert_element_type.1526) + %dot_general.637 = f32[1,128,16,128]{3,1,2,0:T(8,128)} multiply(%convert.387, %convert.388), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/...k,k->...k/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","16","128"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + ROOT %convert.389 = bf16[1,128,16,128]{3,1,2,0:T(8,128)(2,1)S(1)} convert(%dot_general.637) +} + +%fused_computation.54.clone.1 (param_0.1829: bf16[1,128,16,128]) -> (bf16[1,128,16,64], bf16[1,128,16,64]) { + %param_0.1829 = bf16[1,128,16,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %slice.82 = bf16[1,128,16,64]{3,1,2,0:T(8,128)(2,1)} slice(%param_0.1829), slice={[0:1], [0:128], [0:16], [64:128]}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/split" stack_frame_id=0} + %convert.390 = f32[1,128,16,64]{3,1,2,0:T(8,128)} convert(%slice.82) + %neg.134 = f32[1,128,16,64]{3,1,2,0:T(8,128)} negate(%convert.390), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/neg" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","16","64"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %convert.391 = bf16[1,128,16,64]{3,1,2,0:T(8,128)(2,1)S(1)} convert(%neg.134) + %slice.83 = bf16[1,128,16,64]{3,1,2,0:T(8,128)(2,1)S(1)} slice(%param_0.1829), slice={[0:1], [0:128], [0:16], [0:64]}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/split" stack_frame_id=0} + ROOT %tuple.245 = (bf16[1,128,16,64]{3,1,2,0:T(8,128)(2,1)S(1)}, bf16[1,128,16,64]{3,1,2,0:T(8,128)(2,1)S(1)}) tuple(%convert.391, %slice.83) +} + +%fused_computation.58.clone.1 (param_0.1830: bf16[1,128,16,64], param_1.1905: bf16[1,128,16,64], param_2.1454: bf16[128,128], param_3.955: bf16[128,128], param_4.567: f32[128,16], param_5.491: bf16[128,16,128], param_6.344: bf16[128]) -> bf16[16,128,128] { + %param_6.344 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(6) + %dot_general.640 = bf16[1,128,16,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_6.344), dimensions={3}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/...k,k->...k/dot_general" stack_frame_id=0} + %convert.392 = f32[1,128,16,128]{3,1,2,0:T(8,128)} convert(%dot_general.640) + %param_5.491 = bf16[128,16,128]{2,0,1:T(8,128)(2,1)S(1)} parameter(5) + %convert_element_type.1529 = f32[128,16,128]{2,0,1:T(8,128)} convert(%param_5.491), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %bitcast.1024 = f32[1,128,16,128]{3,1,2,0:T(8,128)} bitcast(%convert_element_type.1529), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %param_4.567 = f32[128,16]{0,1:T(8,128)S(1)} parameter(4) + %mul.2402 = f32[1,128,16,128]{3,1,2,0:T(8,128)} broadcast(%param_4.567), dimensions={1,2}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %mul.2401 = f32[1,128,16,128]{3,1,2,0:T(8,128)} multiply(%bitcast.1024, %mul.2402), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %convert_element_type.1528 = bf16[1,128,16,128]{3,1,2,0:T(8,128)(2,1)} convert(%mul.2401), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %convert.393 = f32[1,128,16,128]{3,1,2,0:T(8,128)} convert(%convert_element_type.1528) + %dot_general.639 = f32[1,128,16,128]{3,1,2,0:T(8,128)} multiply(%convert.392, %convert.393), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/...k,k->...k/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","16","128"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %param_3.955 = bf16[128,128]{1,0:T(8,128)(2,1)S(1)} parameter(3) + %mul.2399 = bf16[1,128,16,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_3.955), dimensions={1,3}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %convert.394 = f32[1,128,16,128]{3,1,2,0:T(8,128)} convert(%mul.2399) + %mul.2397 = f32[1,128,16,128]{3,1,2,0:T(8,128)} multiply(%dot_general.639, %convert.394), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","16","128"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %param_1.1905 = bf16[1,128,16,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) + %constant.1136 = bf16[]{:T(256)} constant(-inf) + %pad.85 = bf16[1,128,16,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_1.1905, %constant.1136), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/concatenate" stack_frame_id=0} + %convert.395 = f32[1,128,16,128]{3,1,2,0:T(8,128)} convert(%pad.85) + %param_0.1830 = bf16[1,128,16,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %pad.84 = bf16[1,128,16,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1830, %constant.1136), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/concatenate" stack_frame_id=0} + %convert.396 = f32[1,128,16,128]{3,1,2,0:T(8,128)} convert(%pad.84) + %maximum.61 = f32[1,128,16,128]{3,1,2,0:T(8,128)} maximum(%convert.395, %convert.396), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/concatenate" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","16","128"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %param_2.1454 = bf16[128,128]{1,0:T(8,128)(2,1)S(1)} parameter(2) + %mul.2398 = bf16[1,128,16,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_2.1454), dimensions={1,3}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %convert.397 = f32[1,128,16,128]{3,1,2,0:T(8,128)} convert(%mul.2398) + %mul.2396 = f32[1,128,16,128]{3,1,2,0:T(8,128)} multiply(%maximum.61, %convert.397), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","16","128"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %add.1065 = f32[1,128,16,128]{3,1,2,0:T(8,128)} add(%mul.2397, %mul.2396), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/add" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","16","128"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %constant.425.clone.7 = bf16[]{:T(256)} constant(0.08838) + %mul.2400 = bf16[1,128,16,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%constant.425.clone.7), dimensions={}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %convert.398 = f32[1,128,16,128]{3,1,2,0:T(8,128)} convert(%mul.2400) + %mul.2395 = f32[1,128,16,128]{3,1,2,0:T(8,128)} multiply(%add.1065, %convert.398), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","16","128"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %convert.399 = bf16[1,128,16,128]{3,1,2,0:T(8,128)(2,1)} convert(%mul.2395) + ROOT %bitcast.1023 = bf16[16,128,128]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%convert.399), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/vmap(static_sparsity_splash)/jit(_splash_attention)/splash_mha_fwd_segmented_residuals/splash_mha_fwd_segmented_residuals/squeeze" stack_frame_id=0} +} + +%fused_computation.29.clone.1 (param_0.1815: bf16[4,512,8,128], param_1.1896: s32[]) -> bf16[1,512,8,128] { + %param_0.1815 = bf16[4,512,8,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) + %param_1.1896 = s32[]{:T(128)S(6)} parameter(1) + %constant.420.clone.39 = s32[]{:T(128)} constant(0) + ROOT %dynamic-slice.208 = bf16[1,512,8,128]{1,3,2,0:T(8,128)(2,1)S(1)} dynamic-slice(%param_0.1815, %param_1.1896, %constant.420.clone.39, %constant.420.clone.39, %constant.420.clone.39), dynamic_slice_sizes={1,512,8,128}, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_index_in_dim)/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} +} + +%fused_computation.53.clone.1.clone.clone.clone.1 (param_0.1817: bf16[1,128,2048], param_1.1897: f32[128], param_2.1447: bf16[2048]) -> bf16[128,2048,1] { + %param_2.1447 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(2) + %dot_general.630 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1447), dimensions={2}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/...k,k->...k/dot_general" stack_frame_id=0} + %convert.371 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%dot_general.630) + %param_0.1817 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.1520 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%param_0.1817), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %param_1.1897 = f32[128]{0:T(128)S(1)} parameter(1) + %mul.2382 = f32[1,128,2048]{2,1,0:T(8,128)} broadcast(%param_1.1897), dimensions={1}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %mul.2381 = f32[1,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1520, %mul.2382), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %convert_element_type.1519 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.2381), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %convert.372 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%convert_element_type.1519) + %dot_general.629 = f32[1,128,2048]{2,1,0:T(8,128)} multiply(%convert.371, %convert.372), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/...k,k->...k/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","2048"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %convert.373 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} convert(%dot_general.629) + ROOT %bitcast.1017 = bf16[128,2048,1]{1,0,2:T(8,128)(2,1)} bitcast(%convert.373) +} + +%fused_computation.21.clone.clone.clone.clone.clone.clone.clone.1 (param_0.1816: bf16[1,2048,8,128]) -> bf16[2048,8,128] { + %param_0.1816 = bf16[1,2048,8,128]{1,3,2,0:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.1016 = bf16[2048,8,128]{0,2,1:T(8,128)(2,1)} bitcast(%param_0.1816), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} +} + +%region_4.5 (reduce_sum.152: f32[], reduce_sum.153: f32[]) -> f32[] { + %reduce_sum.152 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/reduce_sum"} + %reduce_sum.153 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/reduce_sum"} + ROOT %reduce_sum.157 = f32[]{:T(128)} add(%reduce_sum.152, %reduce_sum.153), metadata={op_name="reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%fused_computation.64.clone.1 (param_0.1818: bf16[1,2048,8,128], param_1.1898: bf16[1,128,2048], param_2.1448: f32[128], param_3.952: bf16[2048]) -> (f32[128,8], f32[1,128,8,128]) { + %param_1.1898 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %param_2.1448 = f32[128]{0:T(128)S(1)} parameter(2) + %param_3.952 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(3) + %fusion.90.clone.3 = bf16[128,2048,1]{1,0,2:T(8,128)(2,1)} fusion(%param_1.1898, %param_2.1448, %param_3.952), kind=kLoop, calls=%fused_computation.53.clone.1.clone.clone.clone.1 + %param_0.1818 = bf16[1,2048,8,128]{1,3,2,0:T(8,128)(2,1)S(1)} parameter(0) + %fusion.64.clone.3 = bf16[2048,8,128]{0,2,1:T(8,128)(2,1)} fusion(%param_0.1818), kind=kLoop, calls=%fused_computation.21.clone.clone.clone.clone.clone.clone.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} + %convolution.58.clone.3 = bf16[128,8,128]{2,0,1:T(8,128)(2,1)} convolution(%fusion.90.clone.3, %fusion.64.clone.3), window={size=8 pad=7_7 rhs_reversal=1}, dim_labels=bf0_i0o->b0f, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} + %convert_element_type.971.clone.3 = f32[128,8,128]{2,0,1:T(8,128)} convert(%convolution.58.clone.3), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %bitcast.256.clone.3 = f32[1,128,8,128]{3,1,2,0:T(8,128)S(1)} bitcast(%convert_element_type.971.clone.3), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %square.282 = f32[1,128,8,128]{3,1,2,0:T(8,128)} multiply(%bitcast.256.clone.3, %bitcast.256.clone.3), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/square" stack_frame_id=0} + %constant.421.clone.14 = f32[]{:T(128)} constant(0) + %reduce.277 = f32[128,8]{0,1:T(8,128)S(1)} reduce(%square.282, %constant.421.clone.14), dimensions={0,3}, to_apply=%region_4.5, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/reduce_sum" stack_frame_id=0} + ROOT %tuple.242 = (f32[128,8]{0,1:T(8,128)S(1)}, f32[1,128,8,128]{3,1,2,0:T(8,128)S(1)}) tuple(%reduce.277, %bitcast.256.clone.3) +} + +%fused_computation.75.clone.1 (param_0.1819: f32[128,8]) -> f32[128,8] { + %param_0.1819 = f32[128,8]{0,1:T(8,128)S(1)} parameter(0) + %constant.424.clone.9 = f32[]{:T(128)} constant(0.0078125) + %broadcast.1095 = f32[128,8]{0,1:T(8,128)} broadcast(%constant.424.clone.9), dimensions={}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/div" stack_frame_id=0} + %div.1031 = f32[128,8]{0,1:T(8,128)} multiply(%param_0.1819, %broadcast.1095), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/div" stack_frame_id=0} + %constant.423.clone.14 = f32[]{:T(128)} constant(1e-06) + %broadcast.1094 = f32[128,8]{0,1:T(8,128)} broadcast(%constant.423.clone.14), dimensions={}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/add" stack_frame_id=0} + %add.1062 = f32[128,8]{0,1:T(8,128)} add(%div.1031, %broadcast.1094), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/add" stack_frame_id=0} + ROOT %rsqrt.223 = f32[128,8]{0,1:T(8,128)S(1)} rsqrt(%add.1062), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/rsqrt" stack_frame_id=0} +} + +%fused_computation.65.clone.2 (param_0.1820: f32[1,128,8,128], param_1.1899: f32[128,8], param_2.1449: bf16[128]) -> bf16[1,128,8,128] { + %param_2.1449 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(2) + %dot_general.632 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_2.1449), dimensions={3}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/...k,k->...k/dot_general" stack_frame_id=0} + %convert.400 = f32[1,128,8,128]{3,1,2,0:T(8,128)} convert(%dot_general.632) + %param_0.1820 = f32[1,128,8,128]{3,1,2,0:T(8,128)S(1)} parameter(0) + %param_1.1899 = f32[128,8]{0,1:T(8,128)S(1)} parameter(1) + %mul.2384 = f32[1,128,8,128]{3,1,2,0:T(8,128)} broadcast(%param_1.1899), dimensions={1,2}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %mul.2383 = f32[1,128,8,128]{3,1,2,0:T(8,128)} multiply(%param_0.1820, %mul.2384), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %convert_element_type.1521 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)} convert(%mul.2383), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %convert.401 = f32[1,128,8,128]{3,1,2,0:T(8,128)} convert(%convert_element_type.1521) + %dot_general.631 = f32[1,128,8,128]{3,1,2,0:T(8,128)} multiply(%convert.400, %convert.401), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/...k,k->...k/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","8","128"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + ROOT %convert.402 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} convert(%dot_general.631) +} + +%fused_computation.66.clone.1 (param_0.1821: bf16[1,128,8,128]) -> (bf16[1,128,8,64], bf16[1,128,8,64]) { + %param_0.1821 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %slice.80 = bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)} slice(%param_0.1821), slice={[0:1], [0:128], [0:8], [64:128]}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/split" stack_frame_id=0} + %convert.403 = f32[1,128,8,64]{3,1,2,0:T(8,128)} convert(%slice.80) + %neg.133 = f32[1,128,8,64]{3,1,2,0:T(8,128)} negate(%convert.403), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/neg" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","8","64"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %convert.404 = bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} convert(%neg.133) + %slice.81 = bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} slice(%param_0.1821), slice={[0:1], [0:128], [0:8], [0:64]}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/split" stack_frame_id=0} + ROOT %tuple.243 = (bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)}, bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)}) tuple(%convert.404, %slice.81) +} + +%fused_computation.69.clone.1 (param_0.1822: bf16[1,128,8,64], param_1.1900: bf16[1,128,8,64], param_2.1450: bf16[128,128], param_3.953: bf16[128,128], param_4.566: f32[1,128,8,128], param_5.490: f32[128,8], param_6.343: bf16[128]) -> bf16[8,128,128] { + %param_6.343 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(6) + %dot_general.634 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_6.343), dimensions={3}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/...k,k->...k/dot_general" stack_frame_id=0} + %convert.405 = f32[1,128,8,128]{3,1,2,0:T(8,128)} convert(%dot_general.634) + %param_4.566 = f32[1,128,8,128]{3,1,2,0:T(8,128)S(1)} parameter(4) + %param_5.490 = f32[128,8]{0,1:T(8,128)S(1)} parameter(5) + %mul.2390 = f32[1,128,8,128]{3,1,2,0:T(8,128)} broadcast(%param_5.490), dimensions={1,2}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %mul.2389 = f32[1,128,8,128]{3,1,2,0:T(8,128)} multiply(%param_4.566, %mul.2390), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %convert_element_type.1522 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)} convert(%mul.2389), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %convert.406 = f32[1,128,8,128]{3,1,2,0:T(8,128)} convert(%convert_element_type.1522) + %dot_general.633 = f32[1,128,8,128]{3,1,2,0:T(8,128)} multiply(%convert.405, %convert.406), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/...k,k->...k/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","8","128"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %param_3.953 = bf16[128,128]{1,0:T(8,128)(2,1)S(1)} parameter(3) + %mul.2388 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_3.953), dimensions={1,3}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %convert.407 = f32[1,128,8,128]{3,1,2,0:T(8,128)} convert(%mul.2388) + %mul.2386 = f32[1,128,8,128]{3,1,2,0:T(8,128)} multiply(%dot_general.633, %convert.407), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","8","128"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %param_1.1900 = bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) + %constant.1135 = bf16[]{:T(256)} constant(-inf) + %pad.83 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_1.1900, %constant.1135), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/concatenate" stack_frame_id=0} + %convert.408 = f32[1,128,8,128]{3,1,2,0:T(8,128)} convert(%pad.83) + %param_0.1822 = bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %pad.82 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1822, %constant.1135), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/concatenate" stack_frame_id=0} + %convert.409 = f32[1,128,8,128]{3,1,2,0:T(8,128)} convert(%pad.82) + %maximum.60 = f32[1,128,8,128]{3,1,2,0:T(8,128)} maximum(%convert.408, %convert.409), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/concatenate" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","8","128"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %param_2.1450 = bf16[128,128]{1,0:T(8,128)(2,1)S(1)} parameter(2) + %mul.2387 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_2.1450), dimensions={1,3}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %convert.410 = f32[1,128,8,128]{3,1,2,0:T(8,128)} convert(%mul.2387) + %mul.2385 = f32[1,128,8,128]{3,1,2,0:T(8,128)} multiply(%maximum.60, %convert.410), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","8","128"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %add.1063 = f32[1,128,8,128]{3,1,2,0:T(8,128)} add(%mul.2386, %mul.2385), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/add" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","8","128"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %convert.411 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)} convert(%add.1063) + ROOT %bitcast.1018 = bf16[8,128,128]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%convert.411), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/vmap(static_sparsity_splash)/jit(_splash_attention)/splash_mha_fwd_segmented_residuals/splash_mha_fwd_segmented_residuals/squeeze" stack_frame_id=0} +} + +%fused_computation.28.clone.1 (param_0.1810: bf16[4,512,8,128], param_1.1892: s32[]) -> bf16[1,512,8,128] { + %param_0.1810 = bf16[4,512,8,128]{3,1,2,0:T(8,128)(2,1)} parameter(0) + %param_1.1892 = s32[]{:T(128)S(6)} parameter(1) + %constant.420.clone.37 = s32[]{:T(128)} constant(0) + ROOT %dynamic-slice.207 = bf16[1,512,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} dynamic-slice(%param_0.1810, %param_1.1892, %constant.420.clone.37, %constant.420.clone.37, %constant.420.clone.37), dynamic_slice_sizes={1,512,8,128}, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_index_in_dim)/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} +} + +%fused_computation.53.clone.clone.1 (param_0.1812: bf16[1,128,2048], param_1.1893: f32[128], param_2.1444: bf16[2048]) -> bf16[128,2048,1] { + %param_2.1444 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(2) + %dot_general.628 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1444), dimensions={2}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/...k,k->...k/dot_general" stack_frame_id=0} + %convert.368 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%dot_general.628) + %param_0.1812 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.1518 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%param_0.1812), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %param_1.1893 = f32[128]{0:T(128)S(1)} parameter(1) + %mul.2380 = f32[1,128,2048]{2,1,0:T(8,128)} broadcast(%param_1.1893), dimensions={1}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %mul.2379 = f32[1,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1518, %mul.2380), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %convert_element_type.1517 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.2379), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %convert.369 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%convert_element_type.1517) + %dot_general.627 = f32[1,128,2048]{2,1,0:T(8,128)} multiply(%convert.368, %convert.369), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/...k,k->...k/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","2048"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %convert.370 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} convert(%dot_general.627) + ROOT %bitcast.1013 = bf16[128,2048,1]{1,0,2:T(8,128)(2,1)} bitcast(%convert.370) +} + +%fused_computation.19.clone.clone.clone.1 (param_0.1811: bf16[1,2048,8,128]) -> bf16[2048,8,128] { + %param_0.1811 = bf16[1,2048,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.1012 = bf16[2048,8,128]{2,0,1:T(8,128)(2,1)} bitcast(%param_0.1811), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} +} + +%fused_computation.48.clone.1 (param_0.1813: bf16[1,2048,8,128], param_1.1894: bf16[1,128,2048], param_2.1445: f32[128], param_3.951: bf16[2048]) -> bf16[8,128,128] { + %param_1.1894 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %param_2.1445 = f32[128]{0:T(128)S(1)} parameter(2) + %param_3.951 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(3) + %fusion.509 = bf16[128,2048,1]{1,0,2:T(8,128)(2,1)} fusion(%param_1.1894, %param_2.1445, %param_3.951), kind=kLoop, calls=%fused_computation.53.clone.clone.1 + %param_0.1813 = bf16[1,2048,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %fusion.508 = bf16[2048,8,128]{2,0,1:T(8,128)(2,1)} fusion(%param_0.1813), kind=kLoop, calls=%fused_computation.19.clone.clone.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} + %convolution.148 = bf16[128,8,128]{2,0,1:T(8,128)(2,1)} convolution(%fusion.509, %fusion.508), window={size=8 pad=7_7 rhs_reversal=1}, dim_labels=bf0_i0o->b0f, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} + ROOT %bitcast.1014 = bf16[8,128,128]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%convolution.148), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/vmap(static_sparsity_splash)/jit(_splash_attention)/splash_mha_fwd_segmented_residuals/splash_mha_fwd_segmented_residuals/squeeze" stack_frame_id=0} +} + +%fused_computation.59.clone.clone.clone.1 (param_0.1833: bf16[16,128,128]) -> bf16[128,16,128] { + %param_0.1833 = bf16[16,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.1026 = bf16[128,16,128]{2,0,1:T(8,128)(2,1)} bitcast(%param_0.1833) +} + +%fused_computation.15.clone.clone.clone.clone.clone.clone.clone.1 (param_0.1832: bf16[1,16,128,2048]) -> bf16[16,128,2048] { + %param_0.1832 = bf16[1,16,128,2048]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.1025 = bf16[16,128,2048]{2,1,0:T(8,128)(2,1)} bitcast(%param_0.1832), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} +} + +%region_5.7 (reduce_sum.158: f32[], reduce_sum.159: f32[]) -> f32[] { + %reduce_sum.158 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/reduce_sum"} + %reduce_sum.159 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/reduce_sum"} + ROOT %reduce_sum.160 = f32[]{:T(128)} add(%reduce_sum.158, %reduce_sum.159), metadata={op_name="reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%fused_computation.40.clone.1 (param_0.1834: bf16[1,128,2048], param_1.1907: bf16[1,16,128,2048], param_2.1455: bf16[16,128,128]) -> (f32[128], bf16[1,128,2048]) { + %param_0.1834 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert.412 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%param_0.1834) + %param_2.1455 = bf16[16,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %fusion.79.clone.3 = bf16[128,16,128]{2,0,1:T(8,128)(2,1)} fusion(%param_2.1455), kind=kLoop, calls=%fused_computation.59.clone.clone.clone.1 + %param_1.1907 = bf16[1,16,128,2048]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %fusion.47.clone.3 = bf16[16,128,2048]{2,1,0:T(8,128)(2,1)} fusion(%param_1.1907), kind=kLoop, calls=%fused_computation.15.clone.clone.clone.clone.clone.clone.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} + %convolution.50.clone.3 = bf16[128,2048,1]{1,0,2:T(8,128)(2,1)} convolution(%fusion.79.clone.3, %fusion.47.clone.3), window={size=16}, dim_labels=b0f_0io->bf0, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} + %bitcast.236.clone.3 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} bitcast(%convolution.50.clone.3), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} + %convert.413 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%bitcast.236.clone.3) + %add.797.clone.3 = f32[1,128,2048]{2,1,0:T(8,128)} add(%convert.412, %convert.413), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/add" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","2048"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %convert_element_type.1530 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%add.797.clone.3), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %square.284 = f32[1,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1530, %convert_element_type.1530), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/square" stack_frame_id=0} + %constant.421.clone.16 = f32[]{:T(128)} constant(0) + %reduce.279 = f32[128]{0:T(128)S(1)} reduce(%square.284, %constant.421.clone.16), dimensions={0,2}, to_apply=%region_5.7, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/reduce_sum" stack_frame_id=0} + %convert.414 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)S(1)} convert(%add.797.clone.3) + ROOT %tuple.246 = (f32[128]{0:T(128)S(1)}, bf16[1,128,2048]{2,1,0:T(8,128)(2,1)S(1)}) tuple(%reduce.279, %convert.414) +} + +%fused_computation.11.clone.1 (param_0.1841: bf16[4,6144,512], param_1.1912: s32[]) -> bf16[1,6144,512] { + %param_0.1841 = bf16[4,6144,512]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1912 = s32[]{:T(128)S(6)} parameter(1) + %constant.420.clone.44 = s32[]{:T(128)} constant(0) + ROOT %dynamic-slice.213 = bf16[1,6144,512]{2,1,0:T(8,128)(2,1)S(1)} dynamic-slice(%param_0.1841, %param_1.1912, %constant.420.clone.44, %constant.420.clone.44), dynamic_slice_sizes={1,6144,512}, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_index_in_dim)/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} +} + +%fused_computation.12.clone.1 (param_0.1840: bf16[4,512,6144], param_1.1911: s32[]) -> bf16[1,512,6144] { + %param_0.1840 = bf16[4,512,6144]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1911 = s32[]{:T(128)S(6)} parameter(1) + %constant.420.clone.43 = s32[]{:T(128)} constant(0) + ROOT %dynamic-slice.212 = bf16[1,512,6144]{2,1,0:T(8,128)(2,1)S(1)} dynamic-slice(%param_0.1840, %param_1.1911, %constant.420.clone.43, %constant.420.clone.43), dynamic_slice_sizes={1,512,6144}, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_index_in_dim)/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} +} + +%fused_computation.13.clone.1 (param_0.1836: bf16[4,512,6144], param_1.1908: s32[]) -> bf16[1,512,6144] { + %param_0.1836 = bf16[4,512,6144]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1908 = s32[]{:T(128)S(6)} parameter(1) + %constant.420.clone.42 = s32[]{:T(128)} constant(0) + ROOT %dynamic-slice.211 = bf16[1,512,6144]{2,1,0:T(8,128)(2,1)S(1)} dynamic-slice(%param_0.1836, %param_1.1908, %constant.420.clone.42, %constant.420.clone.42), dynamic_slice_sizes={1,512,6144}, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_index_in_dim)/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} +} + +%fused_computation.82.clone.1 (param_0.1835: f32[128]) -> f32[128] { + %param_0.1835 = f32[128]{0:T(128)S(1)} parameter(0) + %constant.422.clone.10 = f32[]{:T(128)} constant(0.00048828125) + %broadcast.1099 = f32[128]{0:T(128)} broadcast(%constant.422.clone.10), dimensions={}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call" stack_frame_id=0} + %div.1033 = f32[128]{0:T(128)} multiply(%param_0.1835, %broadcast.1099), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/div" stack_frame_id=0} + %constant.423.clone.16 = f32[]{:T(128)} constant(1e-06) + %broadcast.1098 = f32[128]{0:T(128)} broadcast(%constant.423.clone.16), dimensions={}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call" stack_frame_id=0} + %add.1066 = f32[128]{0:T(128)} add(%div.1033, %broadcast.1098), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/add" stack_frame_id=0} + ROOT %rsqrt.225 = f32[128]{0:T(128)S(1)} rsqrt(%add.1066), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/rsqrt" stack_frame_id=0} +} + +%fused_computation.51.clone.1.clone.1 (param_0.1838: bf16[1,128,2048], param_1.1909: f32[128], param_2.1456: bf16[2048]) -> bf16[128,2048] { + %param_2.1456 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(2) + %dot_general.642 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1456), dimensions={2}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/...k,k->...k/dot_general" stack_frame_id=0} + %convert.377 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%dot_general.642) + %param_0.1838 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.1532 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%param_0.1838), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %param_1.1909 = f32[128]{0:T(128)S(1)} parameter(1) + %mul.2404 = f32[1,128,2048]{2,1,0:T(8,128)} broadcast(%param_1.1909), dimensions={1}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %mul.2403 = f32[1,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1532, %mul.2404), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %convert_element_type.1531 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.2403), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %convert.378 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%convert_element_type.1531) + %dot_general.641 = f32[1,128,2048]{2,1,0:T(8,128)} multiply(%convert.377, %convert.378), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/...k,k->...k/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","2048"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %convert.379 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} convert(%dot_general.641) + ROOT %bitcast.1028 = bf16[128,2048]{1,0:T(8,128)(2,1)} bitcast(%convert.379) +} + +%fused_computation.10.clone.clone.clone.1 (param_0.1837: bf16[1,2048,6144]) -> bf16[2048,6144] { + %param_0.1837 = bf16[1,2048,6144]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.1027 = bf16[2048,6144]{1,0:T(8,128)(2,1)} bitcast(%param_0.1837), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} +} + +%fused_computation.24.clone.1 (param_0.1839: bf16[1,2048,6144], param_1.1910: bf16[1,128,2048], param_2.1457: f32[128], param_3.956: bf16[2048]) -> bf16[1,128,6144] { + %param_1.1910 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %param_2.1457 = f32[128]{0:T(128)S(1)} parameter(2) + %param_3.956 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(3) + %fusion.511 = bf16[128,2048]{1,0:T(8,128)(2,1)} fusion(%param_1.1910, %param_2.1457, %param_3.956), kind=kLoop, calls=%fused_computation.51.clone.1.clone.1 + %param_0.1839 = bf16[1,2048,6144]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %fusion.510 = bf16[2048,6144]{1,0:T(8,128)(2,1)} fusion(%param_0.1839), kind=kLoop, calls=%fused_computation.10.clone.clone.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} + %convolution.149 = bf16[128,6144]{1,0:T(8,128)(2,1)} convolution(%fusion.511, %fusion.510), dim_labels=bf_io->bf, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} + ROOT %bitcast.1029 = bf16[1,128,6144]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%convolution.149), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} +} + +%fused_computation.51.clone.clone.clone.1 (param_0.1844: bf16[1,128,2048], param_1.1913: f32[128], param_2.1458: bf16[2048]) -> bf16[128,2048] { + %param_2.1458 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(2) + %dot_general.644 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1458), dimensions={2}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/...k,k->...k/dot_general" stack_frame_id=0} + %convert.380 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%dot_general.644) + %param_0.1844 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.1534 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%param_0.1844), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %param_1.1913 = f32[128]{0:T(128)S(1)} parameter(1) + %mul.2406 = f32[1,128,2048]{2,1,0:T(8,128)} broadcast(%param_1.1913), dimensions={1}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %mul.2405 = f32[1,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1534, %mul.2406), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0} + %convert_element_type.1533 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.2405), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %convert.381 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%convert_element_type.1533) + %dot_general.643 = f32[1,128,2048]{2,1,0:T(8,128)} multiply(%convert.380, %convert.381), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/...k,k->...k/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","2048"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %convert.382 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} convert(%dot_general.643) + ROOT %bitcast.1032 = bf16[128,2048]{1,0:T(8,128)(2,1)} bitcast(%convert.382) +} + +%fused_computation.8.clone.clone.clone.clone.clone.clone.clone.clone.clone.1 (param_0.1843: bf16[1,2048,6144]) -> bf16[2048,6144] { + %param_0.1843 = bf16[1,2048,6144]{2,1,0:T(8,128)(2,1)} parameter(0) + ROOT %bitcast.1031 = bf16[2048,6144]{1,0:T(8,128)(2,1)} bitcast(%param_0.1843), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} +} + +%fused_computation.27.clone.clone.clone.1 (param_0.1845: bf16[1,2048,6144], param_1.1914: bf16[1,128,6144], param_2.1459: bf16[1,128,2048], param_3.957: f32[128], param_4.568: bf16[2048]) -> bf16[128,6144] { + %param_1.1914 = bf16[1,128,6144]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %convert.383 = f32[1,128,6144]{2,1,0:T(8,128)} convert(%param_1.1914) + %constant.426.clone.10 = bf16[]{:T(256)} constant(1) + %jit_silu_.46 = bf16[1,128,6144]{2,1,0:T(8,128)(2,1)} broadcast(%constant.426.clone.10), dimensions={}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/jit(silu)" stack_frame_id=0} + %convert.384 = f32[1,128,6144]{2,1,0:T(8,128)} convert(%jit_silu_.46) + %neg.135 = f32[1,128,6144]{2,1,0:T(8,128)} negate(%convert.383), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/jit(silu)/neg" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","6144"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %exp.79 = f32[1,128,6144]{2,1,0:T(8,128)} exponential(%neg.135), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/jit(silu)/exp" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","6144"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %add.1067 = f32[1,128,6144]{2,1,0:T(8,128)} add(%exp.79, %convert.384), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/jit(silu)/add" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","6144"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %div.1034 = f32[1,128,6144]{2,1,0:T(8,128)} divide(%convert.384, %add.1067), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/jit(silu)/div" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","6144"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %mul.2408 = f32[1,128,6144]{2,1,0:T(8,128)} multiply(%convert.383, %div.1034), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/jit(silu)/mul" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","6144"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %param_2.1459 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %param_3.957 = f32[128]{0:T(128)S(1)} parameter(3) + %param_4.568 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(4) + %fusion.513 = bf16[128,2048]{1,0:T(8,128)(2,1)} fusion(%param_2.1459, %param_3.957, %param_4.568), kind=kLoop, calls=%fused_computation.51.clone.clone.clone.1 + %param_0.1845 = bf16[1,2048,6144]{2,1,0:T(8,128)(2,1)} parameter(0) + %fusion.512 = bf16[2048,6144]{1,0:T(8,128)(2,1)} fusion(%param_0.1845), kind=kLoop, calls=%fused_computation.8.clone.clone.clone.clone.clone.clone.clone.clone.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} + %convolution.150 = bf16[128,6144]{1,0:T(8,128)(2,1)} convolution(%fusion.513, %fusion.512), dim_labels=bf_io->bf, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} + %bitcast.1034 = bf16[1,128,6144]{2,1,0:T(8,128)(2,1)} bitcast(%convolution.150), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} + %convert.385 = f32[1,128,6144]{2,1,0:T(8,128)} convert(%bitcast.1034) + %mul.2407 = f32[1,128,6144]{2,1,0:T(8,128)} multiply(%mul.2408, %convert.385), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/mul" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","6144"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %convert.386 = bf16[1,128,6144]{2,1,0:T(8,128)(2,1)} convert(%mul.2407) + ROOT %bitcast.1033 = bf16[128,6144]{1,0:T(8,128)(2,1)} bitcast(%convert.386) +} + +%fused_computation.6.clone.clone.clone.clone.clone.1 (param_0.1842: bf16[1,6144,2048]) -> bf16[6144,2048] { + %param_0.1842 = bf16[1,6144,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.1030 = bf16[6144,2048]{1,0:T(8,128)(2,1)} bitcast(%param_0.1842), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} +} + +%fused_computation.37.clone.1 (param_0.1846: bf16[1,128,2048], param_1.1915: bf16[1,6144,2048], param_2.1460: bf16[1,2048,6144], param_3.958: bf16[1,128,6144], param_4.569: f32[128], param_5.492: bf16[2048]) -> bf16[1,128,2048] { + %param_0.1846 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert.415 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%param_0.1846) + %param_2.1460 = bf16[1,2048,6144]{2,1,0:T(8,128)(2,1)} parameter(2) + %param_3.958 = bf16[1,128,6144]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) + %param_4.569 = f32[128]{0:T(128)S(1)} parameter(4) + %param_5.492 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(5) + %fusion.36.clone.3 = bf16[128,6144]{1,0:T(8,128)(2,1)} fusion(%param_2.1460, %param_3.958, %param_0.1846, %param_4.569, %param_5.492), kind=kOutput, calls=%fused_computation.27.clone.clone.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} + %param_1.1915 = bf16[1,6144,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %fusion.514 = bf16[6144,2048]{1,0:T(8,128)(2,1)} fusion(%param_1.1915), kind=kLoop, calls=%fused_computation.6.clone.clone.clone.clone.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} + %convolution.151 = bf16[128,2048]{1,0:T(8,128)(2,1)} convolution(%fusion.36.clone.3, %fusion.514), dim_labels=bf_io->bf, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} + %bitcast.1035 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} bitcast(%convolution.151), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0} + %convert.416 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%bitcast.1035) + %add.1068 = f32[1,128,2048]{2,1,0:T(8,128)} add(%convert.415, %convert.416), metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/add" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","2048"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + ROOT %convert.417 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)S(1)} convert(%add.1068) +} + +%wide.region_1.8_spmd.sunk.clone.clone.clone.sunk (wide.param.1: (s32[], bf16[1,128,2048], u32[4], u32[4,4], u32[4], /*index=5*/u32[4,4], u32[4], u32[4,4], u32[4], u32[4,4], /*index=10*/u32[4], u32[4,4], u32[4], u32[4,4], bf16[4,1,128,2048], /*index=15*/bf16[4,128], bf16[4,2048], bf16[4,512,16,128], bf16[1,128,1,128], bf16[1,128,1,128], /*index=20*/bf16[4,128], bf16[4,512,8,128], bf16[4,512,8,128], bf16[4,16,128,512], bf16[4,2048], /*index=25*/bf16[4,512,6144], bf16[4,512,6144], bf16[4,6144,512], u32[4], u32[4,4], /*index=30*/u32[4], u32[4,4], u32[4], u32[4,4], u32[4], /*index=35*/u32[4,4], u32[4], u32[4,4], u32[4], u32[4,4], /*index=40*/s32[128], s32[], s8[1,1,1], s8[1,1,1])) -> (s32[], bf16[1,128,2048], u32[4], u32[4,4], u32[4], /*index=5*/u32[4,4], u32[4], u32[4,4], u32[4], u32[4,4], /*index=10*/u32[4], u32[4,4], u32[4], u32[4,4], bf16[4,1,128,2048], /*index=15*/bf16[4,128], bf16[4,2048], bf16[4,512,16,128], bf16[1,128,1,128], bf16[1,128,1,128], /*index=20*/bf16[4,128], bf16[4,512,8,128], bf16[4,512,8,128], bf16[4,16,128,512], bf16[4,2048], /*index=25*/bf16[4,512,6144], bf16[4,512,6144], bf16[4,6144,512], u32[4], u32[4,4], /*index=30*/u32[4], u32[4,4], u32[4], u32[4,4], u32[4], /*index=35*/u32[4,4], u32[4], u32[4,4], u32[4], u32[4,4], /*index=40*/s32[128], s32[], s8[1,1,1], s8[1,1,1]) { + %constant.477.clone..sunk.1 = s32[]{:T(128)} constant(1) + %wide.param.1 = (s32[]{:T(128)}, bf16[1,128,2048]{2,1,0:T(8,128)(2,1)S(1)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, /*index=5*/u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, /*index=10*/u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, bf16[4,1,128,2048]{3,2,1,0:T(8,128)(2,1)}, /*index=15*/bf16[4,128]{1,0:T(4,128)(2,1)S(1)}, bf16[4,2048]{1,0:T(4,128)(2,1)S(1)}, bf16[4,512,16,128]{1,3,2,0:T(8,128)(2,1)S(1)}, bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)S(1)}, bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)S(1)}, /*index=20*/bf16[4,128]{1,0:T(4,128)(2,1)S(1)}, bf16[4,512,8,128]{1,3,2,0:T(8,128)(2,1)}, bf16[4,512,8,128]{3,1,2,0:T(8,128)(2,1)}, bf16[4,16,128,512]{3,2,1,0:T(8,128)(2,1)}, bf16[4,2048]{1,0:T(4,128)(2,1)S(1)}, /*index=25*/bf16[4,512,6144]{2,1,0:T(8,128)(2,1)}, bf16[4,512,6144]{2,1,0:T(8,128)(2,1)}, bf16[4,6144,512]{2,1,0:T(8,128)(2,1)}, u32[4]{0:T(128)S(1)}, u32[4,4]{1,0:T(4,128)S(1)}, /*index=30*/u32[4]{0:T(128)S(1)}, u32[4,4]{1,0:T(4,128)S(1)}, u32[4]{0:T(128)S(1)}, u32[4,4]{1,0:T(4,128)S(1)}, u32[4]{0:T(128)S(1)}, /*index=35*/u32[4,4]{1,0:T(4,128)S(1)}, u32[4]{0:T(128)S(1)}, u32[4,4]{1,0:T(4,128)S(1)}, u32[4]{0:T(128)S(1)}, u32[4,4]{1,0:T(4,128)S(1)}, /*index=40*/s32[128]{0:T(128)S(1)}, s32[]{:T(128)}, s8[1,1,1]{2,1,0:T(4,128)(4,1)}, s8[1,1,1]{2,1,0:T(4,128)(4,1)}) parameter(0) + %get-tuple-element.2262 = s32[]{:T(128)} get-tuple-element(%wide.param.1), index=0 + %copy.255 = s32[]{:T(128)S(6)} copy(%get-tuple-element.2262), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %add.1069 = s32[]{:T(128)} add(%copy.255, %constant.477.clone..sunk.1), metadata={op_name="jit(train_step)/jvp()/while/body/add" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %get-tuple-element.2276 = bf16[4,1,128,2048]{3,2,1,0:T(8,128)(2,1)} get-tuple-element(%wide.param.1), index=14 + %get-tuple-element.2263 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)S(1)} get-tuple-element(%wide.param.1), index=1 + %bitcast_dynamic-update-slice_fusion.5 = bf16[4,1,128,2048]{3,2,1,0:T(8,128)(2,1)} fusion(%get-tuple-element.2276, %copy.255, %get-tuple-element.2263), kind=kLoop, calls=%fused_computation.60.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_update_index_in_dim)/dynamic_update_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[{"indices":["0","3"]}]}} + %get-tuple-element.2329 = bf16[4,16,128,512]{3,2,1,0:T(8,128)(2,1)} get-tuple-element(%wide.param.1), index=23 + %dynamic-slice_convert_fusion.35 = bf16[1,16,128,512]{3,2,1,0:T(8,128)(2,1)S(1)} fusion(%get-tuple-element.2329, %copy.255), kind=kLoop, calls=%fused_computation.22.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %all-gather.108 = bf16[1,16,128,2048]{3,2,1,0:T(8,128)(2,1)S(1)} all-gather(%dynamic-slice_convert_fusion.35), channel_id=105, replica_groups=mesh['axis_0'=1,'axis_1'=1,'axis_2'=4] {'axis_2'}, dimensions={3}, use_global_device_ids=true, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"barrier_config":{"barrier_type":"CUSTOM","id":"0"},"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %get-tuple-element.2348 = s8[1,1,1]{2,1,0:T(4,128)(4,1)} get-tuple-element(%wide.param.1), index=42, metadata={op_name="jit(train_step)/sharding_constraint" stack_frame_id=0} + %get-tuple-element.2349 = s8[1,1,1]{2,1,0:T(4,128)(4,1)} get-tuple-element(%wide.param.1), index=43, metadata={op_name="jit(train_step)/sharding_constraint" stack_frame_id=0} + %get-tuple-element.2323 = bf16[4,512,16,128]{1,3,2,0:T(8,128)(2,1)S(1)} get-tuple-element(%wide.param.1), index=17 + %dynamic-slice_convert_fusion.36 = bf16[1,512,16,128]{1,3,2,0:T(8,128)(2,1)S(1)} fusion(%get-tuple-element.2323, %copy.255), kind=kLoop, calls=%fused_computation.23.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %all-gather.109 = bf16[1,2048,16,128]{1,3,2,0:T(8,128)(2,1)S(1)} all-gather(%dynamic-slice_convert_fusion.36), channel_id=102, replica_groups=mesh['axis_0'=4,'axis_1'=1,'axis_2'=1] {'axis_0'}, dimensions={1}, use_global_device_ids=true, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"barrier_config":{"barrier_type":"CUSTOM","id":"0"},"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %fusion.515 = f32[128]{0:T(128)S(1)} fusion(%get-tuple-element.2263), kind=kLoop, calls=%fused_computation.38.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %add_rsqrt_fusion.8 = f32[128]{0:T(128)S(1)} fusion(%fusion.515), kind=kLoop, calls=%fused_computation.83.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/rsqrt" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %get-tuple-element.2330 = bf16[4,2048]{1,0:T(4,128)(2,1)S(1)} get-tuple-element(%wide.param.1), index=24 + %get-tuple-element.2322 = bf16[4,2048]{1,0:T(4,128)(2,1)S(1)} get-tuple-element(%wide.param.1), index=16 + %convert_reduce_fusion.7 = (bf16[2048]{0:T(1024)(128)(2,1)S(1)}, bf16[2048]{0:T(1024)(128)(2,1)S(1)}) fusion(%get-tuple-element.2330, %copy.255, %get-tuple-element.2322), kind=kLoop, calls=%fused_computation.73.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"fusion_formation_source":{"fusion_creating_compiler_pass":"TpuAdvancedMOF","fusion_id":"42"},"aliasing_operands":{"lists":[]}} + %get-tuple-element.2203 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} get-tuple-element(%convert_reduce_fusion.7), index=1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %fusion.516 = (f32[128,16]{0,1:T(8,128)S(1)}, bf16[128,16,128]{2,0,1:T(8,128)(2,1)S(1)}) fusion(%all-gather.109, %get-tuple-element.2263, %add_rsqrt_fusion.8, %get-tuple-element.2203), kind=kOutput, calls=%fused_computation.44.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %get-tuple-element.2204 = bf16[128,16,128]{2,0,1:T(8,128)(2,1)S(1)} get-tuple-element(%fusion.516), index=1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/reduce_sum" stack_frame_id=0} + %get-tuple-element.2205 = f32[128,16]{0,1:T(8,128)S(1)} get-tuple-element(%fusion.516), index=0, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/reduce_sum" stack_frame_id=0} + %add_rsqrt_fusion.9 = f32[128,16]{0,1:T(8,128)S(1)} fusion(%get-tuple-element.2205), kind=kLoop, calls=%fused_computation.72.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/rsqrt" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %get-tuple-element.2321 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} get-tuple-element(%wide.param.1), index=15 + %get-tuple-element.2326 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} get-tuple-element(%wide.param.1), index=20 + %fusion.517 = (bf16[128]{0:T(256)(128)(2,1)S(1)}, bf16[128]{0:T(256)(128)(2,1)S(1)}) fusion(%get-tuple-element.2321, %copy.255, %get-tuple-element.2326), kind=kLoop, calls=%fused_computation.79.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"fusion_formation_source":{"fusion_creating_compiler_pass":"TpuAdvancedMOF","fusion_id":"42"},"aliasing_operands":{"lists":[]}} + %get-tuple-element.2206 = bf16[128]{0:T(256)(128)(2,1)S(1)} get-tuple-element(%fusion.517), index=0, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %fusion.518 = bf16[1,128,16,128]{3,1,2,0:T(8,128)(2,1)S(1)} fusion(%get-tuple-element.2204, %add_rsqrt_fusion.9, %get-tuple-element.2206), kind=kLoop, calls=%fused_computation.49.clone.2, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/...k,k->...k/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %slice_negate_fusion.12 = (bf16[1,128,16,64]{3,1,2,0:T(8,128)(2,1)S(1)}, bf16[1,128,16,64]{3,1,2,0:T(8,128)(2,1)S(1)}) fusion(%fusion.518), kind=kLoop, calls=%fused_computation.54.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/neg" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %get-tuple-element.2207 = bf16[1,128,16,64]{3,1,2,0:T(8,128)(2,1)S(1)} get-tuple-element(%slice_negate_fusion.12), index=1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/neg" stack_frame_id=0} + %get-tuple-element.2208 = bf16[1,128,16,64]{3,1,2,0:T(8,128)(2,1)S(1)} get-tuple-element(%slice_negate_fusion.12), index=0, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/neg" stack_frame_id=0} + %get-tuple-element.2325 = bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)S(1)} get-tuple-element(%wide.param.1), index=19 + %get-tuple-element.2324 = bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)S(1)} get-tuple-element(%wide.param.1), index=18 + %bitcast.1070 = bf16[128,128]{1,0:T(8,128)(2,1)S(1)} bitcast(%get-tuple-element.2325) + %bitcast.1068 = bf16[128,128]{1,0:T(8,128)(2,1)S(1)} bitcast(%get-tuple-element.2324) + %fusion.519 = bf16[16,128,128]{2,1,0:T(8,128)(2,1)S(1)} fusion(%get-tuple-element.2207, %get-tuple-element.2208, %bitcast.1070, %bitcast.1068, %add_rsqrt_fusion.9, /*index=5*/%get-tuple-element.2204, %get-tuple-element.2206), kind=kLoop, calls=%fused_computation.58.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/vmap(static_sparsity_splash)/jit(_splash_attention)/splash_mha_fwd_segmented_residuals/splash_mha_fwd_segmented_residuals/squeeze" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %get-tuple-element.2327 = bf16[4,512,8,128]{1,3,2,0:T(8,128)(2,1)} get-tuple-element(%wide.param.1), index=21 + %dynamic-slice_convert_fusion.37 = bf16[1,512,8,128]{1,3,2,0:T(8,128)(2,1)S(1)} fusion(%get-tuple-element.2327, %copy.255), kind=kLoop, calls=%fused_computation.29.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %all-gather.110 = bf16[1,2048,8,128]{1,3,2,0:T(8,128)(2,1)S(1)} all-gather(%dynamic-slice_convert_fusion.37), channel_id=103, replica_groups=mesh['axis_0'=4,'axis_1'=1,'axis_2'=1] {'axis_0'}, dimensions={1}, use_global_device_ids=true, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"barrier_config":{"barrier_type":"CUSTOM","id":"0"},"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %multiply_reduce_fusion.48 = (f32[128,8]{0,1:T(8,128)S(1)}, f32[1,128,8,128]{3,1,2,0:T(8,128)S(1)}) fusion(%all-gather.110, %get-tuple-element.2263, %add_rsqrt_fusion.8, %get-tuple-element.2203), kind=kOutput, calls=%fused_computation.64.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %get-tuple-element.2209 = f32[1,128,8,128]{3,1,2,0:T(8,128)S(1)} get-tuple-element(%multiply_reduce_fusion.48), index=1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/reduce_sum" stack_frame_id=0} + %get-tuple-element.2210 = f32[128,8]{0,1:T(8,128)S(1)} get-tuple-element(%multiply_reduce_fusion.48), index=0, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/reduce_sum" stack_frame_id=0} + %add_rsqrt_fusion.10 = f32[128,8]{0,1:T(8,128)S(1)} fusion(%get-tuple-element.2210), kind=kLoop, calls=%fused_computation.75.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/rsqrt" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %get-tuple-element.2211 = bf16[128]{0:T(256)(128)(2,1)S(1)} get-tuple-element(%fusion.517), index=1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %fusion.520 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} fusion(%get-tuple-element.2209, %add_rsqrt_fusion.10, %get-tuple-element.2211), kind=kLoop, calls=%fused_computation.65.clone.2, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/...k,k->...k/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %slice_negate_fusion.13 = (bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)}, bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)}) fusion(%fusion.520), kind=kLoop, calls=%fused_computation.66.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/neg" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %get-tuple-element.2212 = bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} get-tuple-element(%slice_negate_fusion.13), index=1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/neg" stack_frame_id=0} + %get-tuple-element.2213 = bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} get-tuple-element(%slice_negate_fusion.13), index=0, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/neg" stack_frame_id=0} + %bitcast.1071 = bf16[128,128]{1,0:T(8,128)(2,1)S(1)} bitcast(%get-tuple-element.2325) + %bitcast.1069 = bf16[128,128]{1,0:T(8,128)(2,1)S(1)} bitcast(%get-tuple-element.2324) + %fusion.521 = bf16[8,128,128]{2,1,0:T(8,128)(2,1)S(1)} fusion(%get-tuple-element.2212, %get-tuple-element.2213, %bitcast.1071, %bitcast.1069, %get-tuple-element.2209, /*index=5*/%add_rsqrt_fusion.10, %get-tuple-element.2211), kind=kLoop, calls=%fused_computation.69.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/vmap(static_sparsity_splash)/jit(_splash_attention)/splash_mha_fwd_segmented_residuals/splash_mha_fwd_segmented_residuals/squeeze" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %get-tuple-element.2328 = bf16[4,512,8,128]{3,1,2,0:T(8,128)(2,1)} get-tuple-element(%wide.param.1), index=22 + %dynamic-slice_convert_fusion.38 = bf16[1,512,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} fusion(%get-tuple-element.2328, %copy.255), kind=kLoop, calls=%fused_computation.28.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %all-gather.111 = bf16[1,2048,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} all-gather(%dynamic-slice_convert_fusion.38), channel_id=104, replica_groups=mesh['axis_0'=4,'axis_1'=1,'axis_2'=1] {'axis_0'}, dimensions={1}, use_global_device_ids=true, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"barrier_config":{"barrier_type":"CUSTOM","id":"0"},"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %fusion.522 = bf16[8,128,128]{2,1,0:T(8,128)(2,1)S(1)} fusion(%all-gather.111, %get-tuple-element.2263, %add_rsqrt_fusion.8, %get-tuple-element.2203), kind=kOutput, calls=%fused_computation.48.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %get-tuple-element.2346 = s32[128]{0:T(128)S(1)} get-tuple-element(%wide.param.1), index=40 + %squeeze.300 = s32[128,128]{1,0:T(8,128)S(1)} broadcast(%get-tuple-element.2346), dimensions={0}, metadata={op_name="jit(train_step)/shard_map/vmap(static_sparsity_splash)/jit(_splash_attention)/splash_mha_fwd_segmented_residuals/splash_mha_fwd_segmented_residuals/squeeze" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"output_chunk_bound_config":{"output_chunk_bound":["8"]},"aliasing_operands":{"lists":[]}} + %squeeze.301 = s32[8,128]{1,0:T(8,128)S(1)} broadcast(%get-tuple-element.2346), dimensions={1}, metadata={op_name="jit(train_step)/shard_map/vmap(static_sparsity_splash)/jit(_splash_attention)/splash_mha_fwd_segmented_residuals/splash_mha_fwd_segmented_residuals/squeeze" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"output_chunk_bound_config":{"output_chunk_bound":["128"]},"aliasing_operands":{"lists":[]}} + %iota.68 = s32[128,128]{1,0:T(8,128)S(1)} iota(), iota_dimension=0, metadata={op_name="jit(train_step)/shard_map/vmap(static_sparsity_splash)/jit(_splash_attention)/broadcast_in_dim" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %splash_mha_fwd_segmented_residuals.3 = (f32[128,128]{1,0:T(8,128)}, f32[128,128]{1,0:T(8,128)}, f32[128,128]{1,0:T(8,128)}, bf16[16,128,128]{2,1,0:T(8,128)(2,1)S(1)}, f32[16,128,128]{2,1,0:T(8,128)}) custom-call(%get-tuple-element.2348, %get-tuple-element.2349, %fusion.519, %fusion.521, %fusion.522, /*index=5*/%squeeze.300, %squeeze.301, %iota.68), custom_call_target="tpu_custom_call", operand_layout_constraints={s8[1,1,1]{2,1,0}, s8[1,1,1]{2,1,0}, bf16[16,128,128]{2,1,0}, bf16[8,128,128]{2,1,0}, bf16[8,128,128]{2,1,0}, s32[128,128]{1,0}, s32[8,128]{1,0}, s32[128,128]{1,0}}, frontend_attributes={kernel_metadata={ +"xprof_metadata":"{\"block_q\": 128, \"block_kv\": 128, \"block_kv_compute\": 128, \"block_q_dkv\": 128, \"block_kv_dkv\": 128, \"block_kv_dkv_compute\": 128, \"block_q_dq\": 128, \"block_kv_dq\": 128, \"use_fused_bwd_kernel\": false, \"q_layout\": 1, \"k_layout\": 1, \"v_layout\": 1}" +}}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/vmap(static_sparsity_splash)/jit(_splash_attention)/splash_mha_fwd_segmented_residuals/splash_mha_fwd_segmented_residuals/pallas_call" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"custom_call_config":{"body":"TUzvUgFNTElSMjMuMC4wZ2l0AAFVCwEDBQcJAQMLAz8NDxETFRcZGx0fISMlJykrLS8xMzU3OTs9P0FDRUdJA7IECgQzAfUHFxMXFxMPFxcbCwsXCwsLEycPFwsXExNzUwsXGxMXCxMTF6WFCxMXEwsXDwsLCwsLCwsLCwsLCwsXEwsLJwsLCxMLFxcLCwsLFxcLCzMXFxcXEwsPSwsXCwtzCw8LCwsLNxsLGwsbCxsLGwtTGwsjCyMLIwsbCxsLGwURYYGRjWFNKgIqAgH5GxsbGxsbGxsbGw8PEwsLKwsLEwsTCw8TCw8LEwsTCxMLCxMLDxMLFx8LEwsLEwtTFwsTCwsTCw8TCw9TCxMLExMLCx8LEwsvGwsjFwsTCxcPFwsTCxcLEwsfFwsTCxMPFwsTCxcfCxMLHwsTCwsTCwsTCxcfCxMLHy8XDwcFWVkJBV1JATMPBx8jBw8LCysfJxcfIycHHyMrIxsbTy8rAh4aHwMDHzYCHW+eAh26Ar4CAwMfMgIdbzoCFRkZHY4CkgIdTgJSAgMDrgL6AwVLBU0dQgJGAgVPBVEFUx1T8gIDB30lf/YCgYMRDQAdcgJ2Ag0xHdoC3gIdU+oCHXkKAyMLBzEBAAAAAAAAAIAAAAAAAAAAgAAAAAAAAAAjCwUhgAAAAAAAAACAAAAAAAAAAA0vAwMf5gIDAxID/gMdcVoCAwM/YgIFVR1xZgIdd4YCHRoDHgNhZmZpbmVfbWFwPChkMCwgZDEsIGQyKSAtPiAoZDAsIGQxLCBkMik+AGFmZmluZV9tYXA8KGQwLCBkMSkgLT4gKGQwLCBkMSk+AAVXAwM/VwMDP5oCHXmyAgVZAwMfBgMRCwEFWwVdBV8FYQVjBWUFZwVpBWsFbQVvBXEFcwMDH24CHXd+AgV1BXcDB30lf8YCgYMFeQV7BX0jAQEBBX8d+gL+AgMDHxYDBYEFgwWFBYcDAx9uAx2iA6YDBYkFiyMLAxEBAAAAAAAAAB2yA7YDHb4DwgMdygPOAx3eA+IDAwOnqQWNEQs9AxGtrxWxs7W3ubtXF72//8HDBY8BB/v5+Q0tBZEjCwcxEAAAAAAAAAABAAAAAAAAAAEAAAAAAAAABZMRCwkFlQWXBZkFmwEXxcnN0dXb3+Pn6+8DBRvHHTEJWQMFG8sdMQlbAwUbzx0xCV0DBRvTHTMJXwMFG9cd2QlhIwsFIQgAAAAAAAAAgAAAAAAAAAADBRvdHTMJYwMHS/cb4R0zCWUDB0v3G+UdMwlnAwdL9xvpHTMJaQMFG+0dMQlrAwUb8R0xCW0DBRU1F1kjdHB1Lm1lbW9yeV9zcGFjZTx2bWVtPgAjdHB1LnBpcGVsaW5lX21vZGU8c3luY2hyb25vdXM+ACN0cHUuZGltZW5zaW9uX3NlbWFudGljczxhcmJpdHJhcnk+ACN0cHUuZGltZW5zaW9uX3NlbWFudGljczxwYXJhbGxlbD4AI3RwdS5tZW1vcnlfc3BhY2U8c21lbT4AI3RwdS5jb3JlX3R5cGU8dGM+ACN0cHUuZG90X2RpbWVuc2lvbl9udW1iZXJzPFsxXSwgWzFdLCBbMF0sIFswXSwgWzAsIDAsIDEsIDBdLCBbXSwgW10+ACN0cHUuZG90X2RpbWVuc2lvbl9udW1iZXJzPFsxXSwgWzBdLCBbMF0sIFsxXSwgWzAsIDAsIDEsIDFdLCBbXSwgW10+AAMFFTUXWwMFFTUXXQMFFSkXXwMFFSkXYQMFFSkXYwMFFSkXZQMFFSkXZwMFFSkXaQMFFTUXawMFFTUXbREBAREDAR0+AhkFnQWfLUoCCe4lEzImBwWhBaMdVgINBaUdXgIZBacRCxEdagINBakRAQkFqx16AhkFrR2CAhkFrx2KAg0FsQWzHZYCGQW1EQsFHaICDQW3AwMfqgITCRAAAOAPBbkdtgINBbsFvR3CAg0FvyMBCSEBAAAAAQAAAAMAAAAAAAAAHc4C0gIFwR3WAg0FwwXFHeICDQXHEwkBHe4CGQXJHYUNIwEJIQEAAAABAAAAAgAAAAAAAAAFyx0CAxkFzREBAgQdDgMZBc8F0SUFCQAAAAAF0x0iAw0F1QMJiwICjSWPJZElAwMuAzIDBdcjAQMJAQAAAB06Az4DBdkdQgMNBdsDAz9KAxELFR1SA1YDBd0dWgMNBd8dYgNmAwXhHWoDDQXjEwmQzMzMPx12A3oDBeUdfgMZBecdU4YDHYUZHY4DkgMF6R2WAwEF6wMDH54DJRcJAACA/wXtHaoDDQXvAwWXAgSZmwXxHboDDQXzBfUdxgMNBfcF+R3SAw0F+wMDH9oDJRcJAAAAAAX9HeYDDQX/AwWXBgSZmwMJiwYCjSWPJZElAwMf9gMRAQUjYXJpdGguZmFzdG1hdGg8bm9uZT4AI2FyaXRoLm92ZXJmbG93PG5vbmU+ACN2ZWN0b3Iua2luZDxtYXhpbXVtZj4AI3ZlY3Rvci5raW5kPGFkZD4AAQICAycFAgQCBAkX/QcFBQUPRwsBAgQBCQFBF/UHBQIEAgQfRycFAgQCBAEX9QUCBAIECUknAwIECScFAgQCBB8nBwUCBAIEHxf1BQIEAgQBSQcnBQIEAgQNF/UFIQIEAUkX9QcFAgQCBAlHJwcFAgQCBAknBQUCBAEnBQIEBQkFIQEBAQcHERERHSMdFRUVESUBBQsBAQEHBwcBAQEFCwEBAQcHBQEBBN4iBQERAaUHAwExCxEBqwcDU3chAQEBAQEBBwEHAREBEQERAR0BIwEdARUBFQEVAREBJQEDA3UJAwERB0NNAw0FBSEbBhEDAQMjAwMPCQMBEQcPTwMNBSUnHRQPAykJAx9NAwMtNwMJDwYhAwUDUwMDBwMDAwMDBwMDAwUGBwMFBxtXWRMFByMJVRtXWQMDLZMDCQ8GIQMFA10DAwcDAwMDAwcDAwMFBgcDBQcXYWMTBQcjCV8XYWMDAy03AwkPBiEDBQNnAwMHAwMDAwMHAwMDBQYHAwUHGWttEwUHIwlpGWttGQAPAwEFGQAPAwMLAwMDBwYLAwMDAwcGCwMDAwUVBgsDDwkJKy0vFwYRAwEDMQMDOwkDAREHQT0DDQUzNQMDCwMDAwcGCwMDAwMHBgsDAwMFFQYLAw8JBzk7PRcGEQMBAz8bBhEDAQM3AwMPCQMBEQcPTwMNBUNFHRQPA0cJA7PqAgMDhwkDAQMDL1UDASMHLzkDAQVTVQMDBQMDAwMDBQMDAwUGBQMFBxdZWwkGBQMFA10DAwUDAwMDAwUDAwMFBgUDBQcZYWMJBgUDBQNlAwMFAwMDAwMFAwMDAwMFAwMDBQYFAxsJC2lrbQkGBQMZA28DAwUDAwMHBgUDAwNXAwMFAwMDBQYFAxsJDXN1dwkGBQMZA3kDA0WJAwUnB0UmAwMFB3F7fQMDL1UDASMHLzkDAQVBgQMDL1UDASMHLzkDAQVThSkHKzkDAQWDhzcDNgMqAwMTDwYrAxMDiSkHKzkDEwWNiwMDBQMDAwMDBQMDAwUGBQMTBxWRkwkGBQMTA5URB04DRgMDIQWXjwMDBQMDAwcGBQMDA1cFBgUDKQcTm50JBgUDKQOfAwMFAwMDAwMFAwMDBQYFAxMHEaOlCQYFAxMDpw8GQwMTA6ERB0NNAyEFqas5Bl4DAyEFma0DA3IDkwMJDwaCAwMFA7E7BooDAwUHr3+zAwOVmgMDFysHla4DAxcFtbcJBiEDKwO5DwadAwUDuz0HnRMDBQVfvS0HnxMDBQW1vy8HoRMDBQPBAwOj1gMDFysHo+oDAxcFw8UJBiEDKwPHDwYhAwUDyS0HnxMDBQVfvy8HoRMDBQPNHwdREwMFBc9nIQcrEwMFBcvRAwMHAwMDAwMHAwMDBQYHAwUHF9XXEwUHIwm/F9XXAwMHAwMDAwMHAwMDBQYHAwUHGdvdEwUHIwnTGdvdAwMFAwMDBwYFAwMDVwMDBQMDAwUGBQMbCQ/h4+UJBgUDGQPnPwYRAwUD6QMDRYkDBScHRe4DAwUHw+vtAwMFAwMDAwMFAwMDBQYFAwUHG/HzCQYFAwUD9R8HURMDBQXP9yEHKxMDBQX57wMDBwMDAwMDBwMDAwUGBwMFBxv9/xMFByMJ+xv9/wMDh/IDAwEZAA8DAQUZAA8DA3UJAwERB0NNAw0FBUkbBhEDAQNLAwMPCQMBEQcPTwMNBU1PHRQPA1EJA13RAwMFAwMDAwMFAwMDBQYFAwUHGVNVCQYFAwUDVwMDJ6YCAwkPBicDBQNbMQcnEwMFBV1ZAwMFAwMDAwMFAwMDBQYFAwUHG2FjCQYFAwUDZR8HURMDBQVnXzMGEQMZA2kDAwcDAwMDAwcDAwMDAwcDAwMFBgcDGwkdbW9xCQYHAxkDcwkGBwMbA2sTBQd7C3cdbW9xNQfKAhMDBQNZAwMFAwMDAwMFAwMDBQYFAwUHF3t9CQYFAwUDfyEHKxMDBQV5gQMDBwMDAwMDBwMDAwMDBwMDAwUGBwMnCR+Fh4kJBgcDBQOLCQYHAycDgxMFB3sLjx+Fh4kDAy03AwkPBiEDBQORAwMHAwMDAwMHAwMDBQYHAwUHF5WXEwUHIwmTF5WXAwMtNwMJDwYhAwUDmwMDBwMDAwMDBwMDAwUGBwMFBxmfoRMFByMJnRmfoQMDLTcDCQ8GIQMFA6UDAwcDAwMDAwcDAwMFBgcDBQcbqasTBQcjCacbqasZAA8DAQUZAA8NAAELEQHzBwMPDwsBAQEBAQEHAQcBAwMBCQMBAwMBCQMBDQQBBwEDCwsRAQoCBwMrRwsBAQEBAQEHAQcBAwMLAwMDBwYLAwMDAwcGCwMDAwUVBgsDDwkJCw0PFwYRAwEDEQMDOwkDAREHQT0DDQUTFQMDCwMDAwcGCwMDAwMHBgsDAwMFFQYLAw8JBxkbHRcGEQMBAx8DAydzAwElBicDAQUBIwMDAQkDAQMDAQkDAQ0EAQclIScLEQEOAgcDK0cLAQEBAQEBBwEHAQMDCwMDAwcGCwMDAwMHBgsDAwMFFQYLAw8JCQsNDxcGEQMBAxEDAzsJAwERB0E9Aw0FExUDAwsDAwMHBgsDAwMDBwYLAwMDBRUGCwMPCQcZGx0XBhEDAQMfAwMncwMBJQYnAwEFASMDAwEJAwEDAwEJAwENBAEHJSEnCxEBEgIHAw8PCwEBAQEBAQcBBwEDAwEJAwEDAwEJAwENBAEFAwsLEQEWAgcDJz8LAQEBAQEBBwEHAQMDCwMDAwcGCwMDAwMHBgsDAwMFFQYLAw8JCQsNDxcGEQMBAxEDAzsJAwERB0E9Aw0FExUDAwsDAwMHBgsDAwMDBwYLAwMDBRUGCwMPCQcZGx0XBhEDAQMfAwMBCQMBAwMBCQMBDQQBBSMhCxEBGgIHAw8PCwEBAQEBAQcBBwEDAwEJAwEDAwEJAwENBAEFAwsLEQEeAgcDERMLAQEBAQEBBwEHAQMDAQkDAQMDAQkDAQMDAQkDAQ0EAQULDQsRASICBwMREwsBAQEBAQEHAQcBAwMBCQMBAwMBCQMBAwMBCQMBDQQBBQsNCxEBJgIHAxETCwEBAQEBAQcBBwEDAwEJAwEDAwEJAwEDAwEJAwENBAEFCw0LEQEqAgcDDw8LAQEBAQEBBwEHAQMDAQkDAQMDAQkDAQ0EAQcBAwsLEQEuAgcDDw8LAQEBAQEBBwEHAQMDAQkDAQMDAQkDAQ0EAQcBAwsGAwEFAQDmGgICLRkRCxELEQstGSUVCQsRCw0JFQ0XMRsdCQsNIxELEQsVDRETEQsNDQcJCw0HVS0SAiUJHR1HIyEjKS0fCx0nHSVFESkJCwkJCxsZGRkZGRkZGRkZJR0VDR0lEx0XHxsXExMbFxMXLxcXGRcXFw8ZFRkZIxcjGRUlIxkfDw8NCR0RYnVpbHRpbgBzdGFibGVfbW9zYWljAHRwdQBhcml0aAB2ZWN0b3IAbW9kdWxlAGFyaXRoLmNvbnN0YW50AHZlY3Rvci5sb2FkAGFyaXRoLmluZGV4X2Nhc3QAdmVjdG9yLnNoYXBlX2Nhc3QAZnVuYy5mdW5jAGZ1bmMucmV0dXJuAHZlY3Rvci5icm9hZGNhc3QAYXJpdGguY21waQB0cHUudmVjdG9yX3N0b3JlAG1lbXJlZi5sb2FkAGFyaXRoLmV4dHNpAHNjZi55aWVsZABhcml0aC5leHR1aQBzY2YuaWYAYXJpdGgubXVsZgBhcml0aC5hZGRmAGFyaXRoLm11bGkAYXJpdGguZGl2c2kAdHB1Lm1hdG11bABhcml0aC5hZGRpAHZlY3Rvci5tdWx0aV9yZWR1Y3Rpb24AYXJpdGguc3ViZgBtYXRoLmV4cABhcml0aC5kaXZmAGFyaXRoLnRydW5jZgBtYXRoLmxvZwB0cHUuaW90YQBhcml0aC5hbmRpAGFyaXRoLnNlbGVjdABhcml0aC5tYXhpbXVtZgBhcml0aC5leHRmAGZ1bmN0aW9uX3R5cGUAc3ltX25hbWUAdHJhbnNmb3JtX2luZGljZXMAd2luZG93X2JvdW5kcwB2YWx1ZQBwcmVkaWNhdGUAcGlwZWxpbmVfbW9kZQBicm9hZGNhc3RfaW5fZGltOgB0cmFuc2Zvcm1fMAB0cmFuc2Zvcm1fMQB0cmFuc2Zvcm1fMgB0cmFuc2Zvcm1fMwB0cmFuc2Zvcm1fNAB0cmFuc2Zvcm1fNQB0cmFuc2Zvcm1fNgB0cmFuc2Zvcm1fNwB0cmFuc2Zvcm1fOAB0cmFuc2Zvcm1fOQB0cmFuc2Zvcm1fMTAAZ2V0OgBndDoAZXE6AG11bDoAYWRkAG9wZXJhbmRTZWdtZW50U2l6ZXMAc3RyaWRlcwBicm9hZGNhc3RfaW5fZGltL2Jyb2FkY2FzdF9pbl9kaW0AZGltZW5zaW9uX251bWJlcnMAdHJhbnNwb3NlX2xocwB0cmFuc3Bvc2VfbGhzX2hpbnQAdHJhbnNwb3NlX3JocwBraW5kAHJlZHVjdGlvbl9kaW1zAHN0YWJsZV9tb3NhaWMudmVyc2lvbgBkaW1lbnNpb25fc2VtYW50aWNzAGl0ZXJhdGlvbl9ib3VuZHMAc2NhbGFyX3ByZWZldGNoAHNjcmF0Y2hfb3BlcmFuZHMAc3BsYXNoX21oYV9md2Rfc2VnbWVudGVkX3Jlc2lkdWFscwB0cHUuY29yZV90eXBlAHdpbmRvd19wYXJhbXMAZ2V0AF9zcGxhc2hfYXR0ZW50aW9uAC9fX3cvbWF4dGV4dC9tYXh0ZXh0Ly52ZW52L2xpYi9weXRob24zLjEyL3NpdGUtcGFja2FnZXMvamF4L2V4cGVyaW1lbnRhbC9wYWxsYXMvb3BzL3RwdS9zcGxhc2hfYXR0ZW50aW9uL3NwbGFzaF9hdHRlbnRpb25fa2VybmVsLnB5AGNvbnZlcnRfZWxlbWVudF90eXBlOgBjb252ZXJ0X2VsZW1lbnRfdHlwZS9jb252ZXJ0X2VsZW1lbnRfdHlwZQBndABndC9ndABkaXY6AGRpdgBlcQBlcS9lcQBjb25kOgBjb25kAGdldC9nZXQAZmFzdG1hdGgAbXVsL211bABzd2FwOgBzd2FwL3N3YXAAbG9nOgBsb2cvbG9nAGFkZDoAYWRkL2FkZABicm9hZGNhc3RfaW5fZGltAHNjYW46AHNjYW4AbXVsAG92ZXJmbG93RmxhZ3MAZG90X2dlbmVyYWw6AGRvdF9nZW5lcmFsL2RvdF9nZW5lcmFsAGRpbWVuc2lvbnMAaW90YToAaW90YS9pb3RhAGdlOgBnZS9nZQBhbmQ6AGFuZC9hbmQAaml0OgBqaXQAc2VsZWN0X246AHNlbGVjdF9uL3NlbGVjdF9uAHJlZHVjZV9tYXg6AHJlZHVjZV9tYXgvcmVkdWNlX21heABtYXg6AG1heC9tYXgAc3ViOgBzdWIvc3ViAGV4cDoAZXhwL2V4cAByZWR1Y2Vfc3VtOgByZWR1Y2Vfc3VtL3JlZHVjZV9zdW0A","needs_layout_passes":true,"allow_input_fusion":[],"serialization_format":"1","output_memory_colors":[],"output_memory_space_colors":[],"input_memory_space_colors":[]},"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %get-tuple-element.2214 = bf16[16,128,128]{2,1,0:T(8,128)(2,1)S(1)} get-tuple-element(%splash_mha_fwd_segmented_residuals.3), index=3, frontend_attributes={kernel_metadata={ +"xprof_metadata":"{\"block_q\": 128, \"block_kv\": 128, \"block_kv_compute\": 128, \"block_q_dkv\": 128, \"block_kv_dkv\": 128, \"block_kv_dkv_compute\": 128, \"block_q_dq\": 128, \"block_kv_dq\": 128, \"use_fused_bwd_kernel\": false, \"q_layout\": 1, \"k_layout\": 1, \"v_layout\": 1}" +}}, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/shard_map/vmap(static_sparsity_splash)/jit(_splash_attention)/splash_mha_fwd_segmented_residuals/splash_mha_fwd_segmented_residuals/pallas_call" stack_frame_id=0} + %fusion.523 = (f32[128]{0:T(128)S(1)}, bf16[1,128,2048]{2,1,0:T(8,128)(2,1)S(1)}) fusion(%get-tuple-element.2263, %all-gather.108, %get-tuple-element.2214), kind=kOutput, calls=%fused_computation.40.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %get-tuple-element.2215 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)S(1)} get-tuple-element(%fusion.523), index=1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/reduce_sum" stack_frame_id=0} + %get-tuple-element.2333 = bf16[4,6144,512]{2,1,0:T(8,128)(2,1)} get-tuple-element(%wide.param.1), index=27 + %dynamic-slice_convert_fusion.39 = bf16[1,6144,512]{2,1,0:T(8,128)(2,1)S(1)} fusion(%get-tuple-element.2333, %copy.255), kind=kLoop, calls=%fused_computation.11.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %all-gather.112 = bf16[1,6144,2048]{2,1,0:T(8,128)(2,1)S(1)} all-gather(%dynamic-slice_convert_fusion.39), channel_id=108, replica_groups=mesh['axis_0'=1,'axis_1'=4] {'axis_1'}, dimensions={2}, use_global_device_ids=true, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"barrier_config":{"barrier_type":"CUSTOM","id":"0"},"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %get-tuple-element.2332 = bf16[4,512,6144]{2,1,0:T(8,128)(2,1)} get-tuple-element(%wide.param.1), index=26 + %dynamic-slice_convert_fusion.40 = bf16[1,512,6144]{2,1,0:T(8,128)(2,1)S(1)} fusion(%get-tuple-element.2332, %copy.255), kind=kLoop, calls=%fused_computation.12.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %all-gather.113 = bf16[1,2048,6144]{2,1,0:T(8,128)(2,1)} all-gather(%dynamic-slice_convert_fusion.40), channel_id=107, replica_groups=mesh['axis_0'=4,'axis_1'=1] {'axis_0'}, dimensions={1}, use_global_device_ids=true, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"barrier_config":{"barrier_type":"CUSTOM","id":"0"},"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %get-tuple-element.2331 = bf16[4,512,6144]{2,1,0:T(8,128)(2,1)} get-tuple-element(%wide.param.1), index=25 + %dynamic-slice_convert_fusion.41 = bf16[1,512,6144]{2,1,0:T(8,128)(2,1)S(1)} fusion(%get-tuple-element.2331, %copy.255), kind=kLoop, calls=%fused_computation.13.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %all-gather.114 = bf16[1,2048,6144]{2,1,0:T(8,128)(2,1)S(1)} all-gather(%dynamic-slice_convert_fusion.41), channel_id=106, replica_groups=mesh['axis_0'=4,'axis_1'=1] {'axis_0'}, dimensions={1}, use_global_device_ids=true, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"barrier_config":{"barrier_type":"CUSTOM","id":"0"},"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %get-tuple-element.2216 = f32[128]{0:T(128)S(1)} get-tuple-element(%fusion.523), index=0, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/reduce_sum" stack_frame_id=0} + %add_rsqrt_fusion.11 = f32[128]{0:T(128)S(1)} fusion(%get-tuple-element.2216), kind=kLoop, calls=%fused_computation.82.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/rsqrt" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %get-tuple-element.2217 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} get-tuple-element(%convert_reduce_fusion.7), index=0, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/convert_element_type" stack_frame_id=0} + %fusion.524 = bf16[1,128,6144]{2,1,0:T(8,128)(2,1)S(1)} fusion(%all-gather.114, %get-tuple-element.2215, %add_rsqrt_fusion.11, %get-tuple-element.2217), kind=kOutput, calls=%fused_computation.24.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %fusion.525 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)S(1)} fusion(%get-tuple-element.2215, %all-gather.112, %all-gather.113, %fusion.524, %add_rsqrt_fusion.11, /*index=5*/%get-tuple-element.2217), kind=kOutput, calls=%fused_computation.37.clone.1, metadata={op_name="jit(train_step)/jvp()/while/body/closed_call/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %get-tuple-element.2334 = u32[4]{0:T(128)S(1)} get-tuple-element(%wide.param.1), index=28, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_update_index_in_dim)/dynamic_update_slice" stack_frame_id=0} + %copy.272 = u32[4]{0:T(128)} copy(%get-tuple-element.2334), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %get-tuple-element.2335 = u32[4,4]{1,0:T(4,128)S(1)} get-tuple-element(%wide.param.1), index=29, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_update_index_in_dim)/dynamic_update_slice" stack_frame_id=0} + %copy.273 = u32[4,4]{1,0:T(4,128)} copy(%get-tuple-element.2335), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %get-tuple-element.2336 = u32[4]{0:T(128)S(1)} get-tuple-element(%wide.param.1), index=30, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_update_index_in_dim)/dynamic_update_slice" stack_frame_id=0} + %copy.274 = u32[4]{0:T(128)} copy(%get-tuple-element.2336), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %get-tuple-element.2337 = u32[4,4]{1,0:T(4,128)S(1)} get-tuple-element(%wide.param.1), index=31, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_update_index_in_dim)/dynamic_update_slice" stack_frame_id=0} + %copy.275 = u32[4,4]{1,0:T(4,128)} copy(%get-tuple-element.2337), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %get-tuple-element.2338 = u32[4]{0:T(128)S(1)} get-tuple-element(%wide.param.1), index=32, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_update_index_in_dim)/dynamic_update_slice" stack_frame_id=0} + %copy.276 = u32[4]{0:T(128)} copy(%get-tuple-element.2338), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %get-tuple-element.2339 = u32[4,4]{1,0:T(4,128)S(1)} get-tuple-element(%wide.param.1), index=33, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_update_index_in_dim)/dynamic_update_slice" stack_frame_id=0} + %copy.277 = u32[4,4]{1,0:T(4,128)} copy(%get-tuple-element.2339), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %get-tuple-element.2340 = u32[4]{0:T(128)S(1)} get-tuple-element(%wide.param.1), index=34, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_update_index_in_dim)/dynamic_update_slice" stack_frame_id=0} + %copy.278 = u32[4]{0:T(128)} copy(%get-tuple-element.2340), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %get-tuple-element.2341 = u32[4,4]{1,0:T(4,128)S(1)} get-tuple-element(%wide.param.1), index=35, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_update_index_in_dim)/dynamic_update_slice" stack_frame_id=0} + %copy.279 = u32[4,4]{1,0:T(4,128)} copy(%get-tuple-element.2341), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %get-tuple-element.2342 = u32[4]{0:T(128)S(1)} get-tuple-element(%wide.param.1), index=36, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_update_index_in_dim)/dynamic_update_slice" stack_frame_id=0} + %copy.280 = u32[4]{0:T(128)} copy(%get-tuple-element.2342), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %get-tuple-element.2343 = u32[4,4]{1,0:T(4,128)S(1)} get-tuple-element(%wide.param.1), index=37, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_update_index_in_dim)/dynamic_update_slice" stack_frame_id=0} + %copy.281 = u32[4,4]{1,0:T(4,128)} copy(%get-tuple-element.2343), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %get-tuple-element.2344 = u32[4]{0:T(128)S(1)} get-tuple-element(%wide.param.1), index=38, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_update_index_in_dim)/dynamic_update_slice" stack_frame_id=0} + %copy.282 = u32[4]{0:T(128)} copy(%get-tuple-element.2344), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %get-tuple-element.2345 = u32[4,4]{1,0:T(4,128)S(1)} get-tuple-element(%wide.param.1), index=39, metadata={op_name="jit(train_step)/jvp()/while/body/jit(dynamic_update_index_in_dim)/dynamic_update_slice" stack_frame_id=0} + %copy.283 = u32[4,4]{1,0:T(4,128)} copy(%get-tuple-element.2345), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} + %get-tuple-element.2347 = s32[]{:T(128)} get-tuple-element(%wide.param.1), index=41 + ROOT %tuple.250 = (s32[]{:T(128)}, bf16[1,128,2048]{2,1,0:T(8,128)(2,1)S(1)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, /*index=5*/u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, /*index=10*/u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, bf16[4,1,128,2048]{3,2,1,0:T(8,128)(2,1)}, /*index=15*/bf16[4,128]{1,0:T(4,128)(2,1)S(1)}, bf16[4,2048]{1,0:T(4,128)(2,1)S(1)}, bf16[4,512,16,128]{1,3,2,0:T(8,128)(2,1)S(1)}, bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)S(1)}, bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)S(1)}, /*index=20*/bf16[4,128]{1,0:T(4,128)(2,1)S(1)}, bf16[4,512,8,128]{1,3,2,0:T(8,128)(2,1)}, bf16[4,512,8,128]{3,1,2,0:T(8,128)(2,1)}, bf16[4,16,128,512]{3,2,1,0:T(8,128)(2,1)}, bf16[4,2048]{1,0:T(4,128)(2,1)S(1)}, /*index=25*/bf16[4,512,6144]{2,1,0:T(8,128)(2,1)}, bf16[4,512,6144]{2,1,0:T(8,128)(2,1)}, bf16[4,6144,512]{2,1,0:T(8,128)(2,1)}, u32[4]{0:T(128)S(1)}, u32[4,4]{1,0:T(4,128)S(1)}, /*index=30*/u32[4]{0:T(128)S(1)}, u32[4,4]{1,0:T(4,128)S(1)}, u32[4]{0:T(128)S(1)}, u32[4,4]{1,0:T(4,128)S(1)}, u32[4]{0:T(128)S(1)}, /*index=35*/u32[4,4]{1,0:T(4,128)S(1)}, u32[4]{0:T(128)S(1)}, u32[4,4]{1,0:T(4,128)S(1)}, u32[4]{0:T(128)S(1)}, u32[4,4]{1,0:T(4,128)S(1)}, /*index=40*/s32[128]{0:T(128)S(1)}, s32[]{:T(128)}, s8[1,1,1]{2,1,0:T(4,128)(4,1)}, s8[1,1,1]{2,1,0:T(4,128)(4,1)}) tuple(%add.1069, %fusion.525, %copy.272, %copy.273, %copy.274, /*index=5*/%copy.275, %copy.276, %copy.277, %copy.278, %copy.279, /*index=10*/%copy.280, %copy.281, %copy.282, %copy.283, %bitcast_dynamic-update-slice_fusion.5, /*index=15*/%get-tuple-element.2321, %get-tuple-element.2322, %get-tuple-element.2323, %get-tuple-element.2324, %get-tuple-element.2325, /*index=20*/%get-tuple-element.2326, %get-tuple-element.2327, %get-tuple-element.2328, %get-tuple-element.2329, %get-tuple-element.2330, /*index=25*/%get-tuple-element.2331, %get-tuple-element.2332, %get-tuple-element.2333, %get-tuple-element.2334, %get-tuple-element.2335, /*index=30*/%get-tuple-element.2336, %get-tuple-element.2337, %get-tuple-element.2338, %get-tuple-element.2339, %get-tuple-element.2340, /*index=35*/%get-tuple-element.2341, %get-tuple-element.2342, %get-tuple-element.2343, %get-tuple-element.2344, %get-tuple-element.2345, /*index=40*/%get-tuple-element.2346, %get-tuple-element.2347, %get-tuple-element.2348, %get-tuple-element.2349) +} + +%wide.region_6.9_spmd.clone.clone.clone (wide.param.113: (s32[], bf16[1,128,2048], u32[4], u32[4,4], u32[4], /*index=5*/u32[4,4], u32[4], u32[4,4], u32[4], u32[4,4], /*index=10*/u32[4], u32[4,4], u32[4], u32[4,4], bf16[4,1,128,2048], /*index=15*/bf16[4,128], bf16[4,2048], bf16[4,512,16,128], bf16[1,128,1,128], bf16[1,128,1,128], /*index=20*/bf16[4,128], bf16[4,512,8,128], bf16[4,512,8,128], bf16[4,16,128,512], bf16[4,2048], /*index=25*/bf16[4,512,6144], bf16[4,512,6144], bf16[4,6144,512], u32[4], u32[4,4], /*index=30*/u32[4], u32[4,4], u32[4], u32[4,4], u32[4], /*index=35*/u32[4,4], u32[4], u32[4,4], u32[4], u32[4,4], /*index=40*/s32[128], s32[], s8[1,1,1], s8[1,1,1])) -> pred[] { + %constant.427.clone.4 = s32[]{:T(128)} constant(4) + %wide.param.113 = (s32[]{:T(128)}, bf16[1,128,2048]{2,1,0:T(8,128)(2,1)S(1)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, /*index=5*/u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, /*index=10*/u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, u32[4]{0:T(128)}, u32[4,4]{1,0:T(4,128)}, bf16[4,1,128,2048]{3,2,1,0:T(8,128)(2,1)}, /*index=15*/bf16[4,128]{1,0:T(4,128)(2,1)S(1)}, bf16[4,2048]{1,0:T(4,128)(2,1)S(1)}, bf16[4,512,16,128]{1,3,2,0:T(8,128)(2,1)S(1)}, bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)S(1)}, bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)S(1)}, /*index=20*/bf16[4,128]{1,0:T(4,128)(2,1)S(1)}, bf16[4,512,8,128]{1,3,2,0:T(8,128)(2,1)}, bf16[4,512,8,128]{3,1,2,0:T(8,128)(2,1)}, bf16[4,16,128,512]{3,2,1,0:T(8,128)(2,1)}, bf16[4,2048]{1,0:T(4,128)(2,1)S(1)}, /*index=25*/bf16[4,512,6144]{2,1,0:T(8,128)(2,1)}, bf16[4,512,6144]{2,1,0:T(8,128)(2,1)}, bf16[4,6144,512]{2,1,0:T(8,128)(2,1)}, u32[4]{0:T(128)S(1)}, u32[4,4]{1,0:T(4,128)S(1)}, /*index=30*/u32[4]{0:T(128)S(1)}, u32[4,4]{1,0:T(4,128)S(1)}, u32[4]{0:T(128)S(1)}, u32[4,4]{1,0:T(4,128)S(1)}, u32[4]{0:T(128)S(1)}, /*index=35*/u32[4,4]{1,0:T(4,128)S(1)}, u32[4]{0:T(128)S(1)}, u32[4,4]{1,0:T(4,128)S(1)}, u32[4]{0:T(128)S(1)}, u32[4,4]{1,0:T(4,128)S(1)}, /*index=40*/s32[128]{0:T(128)S(1)}, s32[]{:T(128)}, s8[1,1,1]{2,1,0:T(4,128)(4,1)}, s8[1,1,1]{2,1,0:T(4,128)(4,1)}) parameter(0) + %get-tuple-element.1920 = s32[]{:T(128)} get-tuple-element(%wide.param.113), index=0 + ROOT %lt.34 = pred[]{:T(512)} compare(%get-tuple-element.1920, %constant.427.clone.4), direction=LT, metadata={op_name="jit(train_step)/jvp()/while/cond/lt" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.425 (param_0.1220: pred[4,128]) -> s32[] { - %param_0.1220 = pred[4,128]{1,0:T(4,128)(4,1)S(1)} parameter(0) - %convert_element_type.1403 = s32[4,128]{1,0:T(4,128)} convert(%param_0.1220), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} - %constant.1066 = s32[]{:T(128)} constant(0) - ROOT %reduce.220 = s32[]{:T(128)} reduce(%convert_element_type.1403, %constant.1066), dimensions={0,1}, to_apply=%region_0.1, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} +%region_7.10 (reduce_sum.171: f32[], reduce_sum.184: f32[]) -> f32[] { + %reduce_sum.171 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} + %reduce_sum.184 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} + ROOT %reduce_sum.185 = f32[]{:T(128)} add(%reduce_sum.171, %reduce_sum.184), metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%fused_computation.406 (param_0.1550: bf16[1,128,2048]) -> f32[128] { + %param_0.1550 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.1283 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%param_0.1550), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} + %square.243 = f32[1,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1283, %convert_element_type.1283), metadata={op_name="jit(train_step)/jvp()/square" stack_frame_id=0} + %constant.460.clone.29 = f32[]{:T(128)} constant(0) + ROOT %reduce.222 = f32[128]{0:T(128)S(1)} reduce(%square.243, %constant.460.clone.29), dimensions={0,2}, to_apply=%region_7.10, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} +} + +%fused_computation.467 (param_0.1364: f32[128]) -> f32[128] { + %param_0.1364 = f32[128]{0:T(128)S(1)} parameter(0) + %constant.469.clone.2 = f32[]{:T(128)} constant(0.00048828125) + %broadcast.972 = f32[128]{0:T(128)} broadcast(%constant.469.clone.2), dimensions={}, metadata={op_name="broadcast.418"} + %div.784 = f32[128]{0:T(128)} multiply(%param_0.1364, %broadcast.972), metadata={op_name="jit(train_step)/jvp()/div" stack_frame_id=0} + %constant.470.clone.2 = f32[]{:T(128)} constant(1e-06) + %broadcast.970 = f32[128]{0:T(128)} broadcast(%constant.470.clone.2), dimensions={}, metadata={op_name="jit(train_step)/jvp()/add" stack_frame_id=0} + %add.940 = f32[128]{0:T(128)} add(%div.784, %broadcast.970), metadata={op_name="jit(train_step)/jvp()/add" stack_frame_id=0} + %bitcast.624 = f32[1,128]{1,0:T(1,128)} bitcast(%add.940), metadata={op_name="jit(train_step)/jvp()/add" stack_frame_id=0} + %rsqrt.195 = f32[1,128]{1,0:T(1,128)} rsqrt(%bitcast.624), metadata={op_name="jit(train_step)/jvp()/rsqrt" stack_frame_id=0} + ROOT %bitcast.614 = f32[128]{0:T(128)} bitcast(%rsqrt.195), metadata={op_name="jit(train_step)/jvp()/rsqrt" stack_frame_id=0} +} + +%fused_computation.412.clone.clone (param_0.1510: bf16[2048], param_1.1694: f32[128], param_2.1284: bf16[1,128,2048]) -> bf16[128,2048] { + %param_0.1510 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(0) + %dot_general.506 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_0.1510), dimensions={2}, metadata={op_name="jit(train_step)/jvp(...k,k->...k)/dot_general" stack_frame_id=0} + %convert.210 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%dot_general.506) + %param_2.1284 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} parameter(2) + %convert_element_type.1364 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%param_2.1284), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} + %param_1.1694 = f32[128]{0:T(128)S(1)} parameter(1) + %mul.2124 = f32[1,128,2048]{2,1,0:T(8,128)} broadcast(%param_1.1694), dimensions={1}, metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + %mul.2123 = f32[1,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1364, %mul.2124), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + %convert_element_type.1363 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.2123), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} + %convert.211 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%convert_element_type.1363) + %dot_general.505 = f32[1,128,2048]{2,1,0:T(8,128)} multiply(%convert.210, %convert.211), metadata={op_name="jit(train_step)/jvp(...k,k->...k)/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","2048"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %convert.212 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} convert(%dot_general.505) + ROOT %bitcast.749 = bf16[128,2048]{1,0:T(8,128)(2,1)} bitcast(%convert.212) +} + +%bitcast_fusion.3 (bitcast_input.3: bf16[151936,2048]) -> bf16[151936,2048] { + %bitcast_input.3 = bf16[151936,2048]{1,0:T(8,128)(2,1)} parameter(0) + ROOT %bitcast.761 = bf16[151936,2048]{1,0:T(8,128)(2,1)} bitcast(%bitcast_input.3) } -%fused_computation.428 (param_0.1203: f32[4,128]) -> f32[4,128] { - %param_0.1203 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %constant.1046 = f32[]{:T(128)} constant(0.00048828125) - %broadcast.829 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1046), dimensions={}, metadata={op_name="broadcast.399"} - %div.759 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.1203, %broadcast.829), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/div" stack_frame_id=0} - %constant.1044 = f32[]{:T(128)} constant(1e-06) - %add.924 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1044), dimensions={}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/add" stack_frame_id=0} - %add.921 = f32[4,128]{1,0:T(4,128)} add(%div.759, %add.924), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/add" stack_frame_id=0} - ROOT %rsqrt.166 = f32[4,128]{1,0:T(4,128)S(1)} rsqrt(%add.921), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/rsqrt" stack_frame_id=0} +%region_8.11 (reduce_max.6: bf16[], reduce_max.8: bf16[]) -> bf16[] { + %reduce_max.6 = bf16[]{:T(256)} parameter(0), metadata={op_name="jit(train_step)/jvp()/reduce_max"} + %reduce_max.8 = bf16[]{:T(256)} parameter(1), metadata={op_name="jit(train_step)/jvp()/reduce_max"} + ROOT %reduce_max.9 = bf16[]{:T(256)} maximum(%reduce_max.6, %reduce_max.8), metadata={op_name="jit(train_step)/jvp()/reduce_max" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.429 (param_0.1204: pred[4,128], param_1.1575: f32[]) -> f32[4,128] { - %param_0.1204 = pred[4,128]{1,0:T(4,128)(4,1)S(1)} parameter(0) - %param_1.1575 = f32[]{:T(128)S(6)} parameter(1) - %broadcast_in_dim.288 = f32[4,128]{1,0:T(4,128)} broadcast(%param_1.1575), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/broadcast_in_dim" stack_frame_id=0} - %constant.1223 = f32[]{:T(128)} constant(0) - %broadcast.833 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1223), dimensions={}, metadata={op_name="broadcast.99"} - ROOT %mul.1920 = f32[4,128]{1,0:T(4,128)S(1)} select(%param_0.1204, %broadcast_in_dim.288, %broadcast.833), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} +%fused_computation.496 (param_0.1518: bf16[151936,2048], param_1.1699: bf16[2048], param_2.1301: f32[128], param_3.867: bf16[1,128,2048]) -> (bf16[128], bf16[128,151936]) { + %param_1.1699 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(1) + %param_2.1301 = f32[128]{0:T(128)S(1)} parameter(2) + %param_3.867 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} parameter(3) + %fusion.319.clone.1 = bf16[128,2048]{1,0:T(8,128)(2,1)} fusion(%param_1.1699, %param_2.1301, %param_3.867), kind=kLoop, calls=%fused_computation.412.clone.clone + %param_0.1518 = bf16[151936,2048]{1,0:T(8,128)(2,1)} parameter(0) + %fusion.357 = bf16[151936,2048]{1,0:T(8,128)(2,1)} fusion(%param_0.1518), kind=kLoop, calls=%bitcast_fusion.3 + %convolution.115.clone.1 = bf16[128,151936]{1,0:T(8,128)(2,1)S(1)} convolution(%fusion.319.clone.1, %fusion.357), dim_labels=bf_oi->bf, metadata={op_name="jit(train_step)/jvp()/dot_general" stack_frame_id=0} + %constant.471.clone.3 = bf16[]{:T(256)} constant(-inf) + %reduce.243 = bf16[128]{0:T(256)(128)(2,1)S(1)} reduce(%convolution.115.clone.1, %constant.471.clone.3), dimensions={1}, to_apply=%region_8.11, metadata={op_name="jit(train_step)/jvp()/reduce_max" stack_frame_id=0} + ROOT %tuple.171 = (bf16[128]{0:T(256)(128)(2,1)S(1)}, bf16[128,151936]{1,0:T(8,128)(2,1)S(1)}) tuple(%reduce.243, %convolution.115.clone.1) } -%fused_computation.431 () -> f32[64] { - %constant.1049 = f32[]{:T(128)} constant(1e+06) - %broadcast.840 = f32[64]{0:T(128)} broadcast(%constant.1049), dimensions={}, metadata={op_name="broadcast.390"} - %iota.46 = s32[64]{0:T(128)} iota(), iota_dimension=0, metadata={op_name="jit(train_step)/layers/iota" stack_frame_id=0} - %constant.1048 = s32[]{:T(128)} constant(2) - %broadcast.839 = s32[64]{0:T(128)} broadcast(%constant.1048), dimensions={}, metadata={op_name="broadcast.391"} - %mul.1921 = s32[64]{0:T(128)} multiply(%iota.46, %broadcast.839), metadata={op_name="jit(train_step)/layers/mul" stack_frame_id=0} - %convert_element_type.1404 = f32[64]{0:T(128)} convert(%mul.1921), metadata={op_name="jit(train_step)/layers/convert_element_type" stack_frame_id=0} - %constant.1047 = f32[]{:T(128)} constant(0.0078125) - %broadcast.838 = f32[64]{0:T(128)} broadcast(%constant.1047), dimensions={}, metadata={op_name="broadcast.392"} - %div.768 = f32[64]{0:T(128)} multiply(%convert_element_type.1404, %broadcast.838), metadata={op_name="jit(train_step)/layers/div" stack_frame_id=0} - ROOT %pow.36 = f32[64]{0:T(128)S(1)} power(%broadcast.840, %div.768), metadata={op_name="jit(train_step)/layers/pow" stack_frame_id=0} +%region_9.12 (reduce_sum.186: f32[], reduce_sum.190: f32[]) -> f32[] { + %reduce_sum.186 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} + %reduce_sum.190 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} + ROOT %reduce_sum.191 = f32[]{:T(128)} add(%reduce_sum.186, %reduce_sum.190), metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.432 (param_0.1218: s32[4,128]) -> (f32[4,128,1,1], f32[4,128]) { - %param_0.1218 = s32[4,128]{1,0:T(4,128)} parameter(0) - %convert_element_type.1405 = f32[4,128]{1,0:T(4,128)S(1)} convert(%param_0.1218), metadata={op_name="jit(train_step)/layers/convert_element_type" stack_frame_id=0} - %bitcast.418 = f32[4,128,1,1]{1,0,3,2:T(4,128)} bitcast(%convert_element_type.1405), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %tuple.162 = (f32[4,128,1,1]{1,0,3,2:T(4,128)}, f32[4,128]{1,0:T(4,128)S(1)}) tuple(%bitcast.418, %convert_element_type.1405) +%fused_computation.328 (param_0.1549: bf16[128,151936], param_1.1725: bf16[128]) -> f32[128] { + %param_0.1549 = bf16[128,151936]{1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.1250 = f32[128,151936]{1,0:T(8,128)} convert(%param_0.1549), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} + %bitcast.520 = f32[1,128,151936]{2,1,0:T(8,128)} bitcast(%convert_element_type.1250), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} + %param_1.1725 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(1) + %sub.100 = f32[1,128,151936]{2,1,0:T(8,128)} broadcast(%param_1.1725), dimensions={1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %sub.71 = f32[1,128,151936]{2,1,0:T(8,128)} subtract(%bitcast.520, %sub.100), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %exp.51 = f32[1,128,151936]{2,1,0:T(8,128)} exponential(%sub.71), metadata={op_name="jit(train_step)/jvp()/exp" stack_frame_id=0} + %constant.460.clone.28 = f32[]{:T(128)} constant(0) + ROOT %reduce.199 = f32[128]{0:T(128)S(1)} reduce(%exp.51, %constant.460.clone.28), dimensions={0,2}, to_apply=%region_9.12, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} +} + +%fused_computation.495 (param_0.1506: f32[128], param_1.1724: bf16[128]) -> (f32[128], f32[128]) { + %param_0.1506 = f32[128]{0:T(128)S(1)} parameter(0) + %log.23 = f32[128]{0:T(128)S(1)} log(%param_0.1506), metadata={op_name="jit(train_step)/jvp()/log" stack_frame_id=0} + %param_1.1724 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(1) + %reduce_max.18.clone.1 = f32[128]{0:T(128)} convert(%param_1.1724), metadata={op_name="jit(train_step)/jvp()/reduce_max" stack_frame_id=0} + %add.945.clone.1 = f32[128]{0:T(128)} add(%log.23, %reduce_max.18.clone.1), metadata={op_name="jit(train_step)/jvp()/add" stack_frame_id=0} + %constant.460.clone.27 = f32[]{:T(128)} constant(0) + %broadcast.968.clone.1 = f32[128]{0:T(128)} broadcast(%constant.460.clone.27), dimensions={}, metadata={op_name="broadcast.120"} + %mul.1958.clone.1 = f32[128]{0:T(128)} multiply(%add.945.clone.1, %broadcast.968.clone.1), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %constant.448.clone.1.clone.1 = f32[]{:T(128)} constant(1) + %broadcast.965.clone.1 = f32[128]{0:T(128)} broadcast(%constant.448.clone.1.clone.1), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/add" stack_frame_id=0} + %add.934.clone.1 = f32[128]{0:T(128)S(1)} add(%mul.1958.clone.1, %broadcast.965.clone.1), metadata={op_name="jit(train_step)/transpose(jvp())/add" stack_frame_id=0} + ROOT %tuple.169 = (f32[128]{0:T(128)S(1)}, f32[128]{0:T(128)S(1)}) tuple(%log.23, %add.934.clone.1) +} + +%region_54.59 (reduce_sum.389: f32[], reduce_sum.393: f32[]) -> f32[] { + %reduce_sum.389 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} + %reduce_sum.393 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} + ROOT %reduce_sum.394 = f32[]{:T(128)} add(%reduce_sum.389, %reduce_sum.393), metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.435 (param_0.1360: f32[2048,4]) -> bf16[4,2048] { - %param_0.1360 = f32[2048,4]{0,1:T(4,128)} parameter(0) - %bitcast.531 = f32[4,2048]{1,0:T(4,128)} bitcast(%param_0.1360), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - ROOT %convert.145 = bf16[4,2048]{1,0:T(4,128)(2,1)} convert(%bitcast.531) +%fused_computation.334 (param_0.1548: bf16[128,151936], param_1.1723: f32[128], param_2.1326: bf16[128], param_3.886: s32[128]) -> f32[128] { + %param_3.886 = s32[128]{0:T(128)S(1)} parameter(3) + %eq.26 = s32[1,128,151936]{2,1,0:T(8,128)} broadcast(%param_3.886), dimensions={1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} + %iota.51 = s32[1,128,151936]{2,1,0:T(8,128)} iota(), iota_dimension=2, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} + %eq.21 = pred[1,128,151936]{2,1,0:T(8,128)(4,1)} compare(%eq.26, %iota.51), direction=EQ, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} + %param_0.1548 = bf16[128,151936]{1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.1263 = f32[128,151936]{1,0:T(8,128)} convert(%param_0.1548), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} + %bitcast.532 = f32[1,128,151936]{2,1,0:T(8,128)} bitcast(%convert_element_type.1263), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} + %param_2.1326 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(2) + %sub.99 = f32[1,128,151936]{2,1,0:T(8,128)} broadcast(%param_2.1326), dimensions={1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %sub.92 = f32[1,128,151936]{2,1,0:T(8,128)} subtract(%bitcast.532, %sub.99), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %param_1.1723 = f32[128]{0:T(128)S(1)} parameter(1) + %sub.97 = f32[1,128,151936]{2,1,0:T(8,128)} broadcast(%param_1.1723), dimensions={1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %sub.91 = f32[1,128,151936]{2,1,0:T(8,128)} subtract(%sub.92, %sub.97), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %constant.460.clone.26 = f32[]{:T(128)} constant(0) + %broadcast.906 = f32[1,128,151936]{2,1,0:T(8,128)} broadcast(%constant.460.clone.26), dimensions={}, metadata={op_name="broadcast.130"} + %mul.1808 = f32[1,128,151936]{2,1,0:T(8,128)} select(%eq.21, %sub.91, %broadcast.906), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + ROOT %reduce.200 = f32[128]{0:T(128)S(1)} reduce(%mul.1808, %constant.460.clone.26), dimensions={0,2}, to_apply=%region_54.59, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} } -%fused_computation.436 (param_0.1359: f32[2048,4]) -> bf16[4,2048] { - %param_0.1359 = f32[2048,4]{0,1:T(4,128)} parameter(0) - %bitcast.530 = f32[4,2048]{1,0:T(4,128)} bitcast(%param_0.1359), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - ROOT %convert.147 = bf16[4,2048]{1,0:T(4,128)(2,1)} convert(%bitcast.530) +%region_0.1 (reduce_sum.137: s32[], reduce_sum.138: s32[]) -> s32[] { + %reduce_sum.137 = s32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} + %reduce_sum.138 = s32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} + ROOT %reduce_sum.139 = s32[]{:T(128)} add(%reduce_sum.137, %reduce_sum.138), metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.437 (param_0.1361: f32[128,4]) -> bf16[4,128] { - %param_0.1361 = f32[128,4]{0,1:T(4,128)} parameter(0) - %bitcast.532 = f32[4,128]{1,0:T(4,128)} bitcast(%param_0.1361), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - ROOT %convert.149 = bf16[4,128]{1,0:T(4,128)(2,1)} convert(%bitcast.532) +%fused_computation.472 (param_0.1369: s32[1,128]) -> s32[] { + %param_0.1369 = s32[1,128]{1,0:T(1,128)S(1)} parameter(0) + %constant.449.clone.8 = s32[]{:T(128)} constant(0) + %broadcast.983 = s32[1,128]{1,0:T(1,128)} broadcast(%constant.449.clone.8), dimensions={}, metadata={op_name="broadcast.99"} + %ne.10 = pred[1,128]{1,0:T(4,128)(4,1)} compare(%param_0.1369, %broadcast.983), direction=NE, metadata={op_name="jit(train_step)/jvp()/ne" stack_frame_id=0} + %convert_element_type.1323 = s32[1,128]{1,0:T(1,128)} convert(%ne.10), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} + ROOT %reduce.240 = s32[]{:T(128)} reduce(%convert_element_type.1323, %constant.449.clone.8), dimensions={0,1}, to_apply=%region_0.1, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} } -%fused_computation.438 (param_0.1362: f32[128,4]) -> bf16[4,128] { - %param_0.1362 = f32[128,4]{0,1:T(4,128)} parameter(0) - %bitcast.533 = f32[4,128]{1,0:T(4,128)} bitcast(%param_0.1362), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - ROOT %convert.151 = bf16[4,128]{1,0:T(4,128)(2,1)} convert(%bitcast.533) +%region_0.1.clone (reduce_sum.612: s32[], reduce_sum.613: s32[]) -> s32[] { + %reduce_sum.612 = s32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} + %reduce_sum.613 = s32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} + ROOT %reduce_sum.614 = s32[]{:T(128)} add(%reduce_sum.612, %reduce_sum.613), metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%region_8.11 (reduce_max.6: bf16[], reduce_max.8: bf16[]) -> bf16[] { - %reduce_max.8 = bf16[]{:T(256)} parameter(1), metadata={op_name="jit(train_step)/jvp()/reduce_max"} - %reduce_max.6 = bf16[]{:T(256)} parameter(0), metadata={op_name="jit(train_step)/jvp()/reduce_max"} - ROOT %reduce_max.9 = bf16[]{:T(256)} maximum(%reduce_max.6, %reduce_max.8), metadata={op_name="jit(train_step)/jvp()/reduce_max" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +%region_69.74 (reduce_sum.470: f32[], reduce_sum.471: f32[]) -> f32[] { + %reduce_sum.470 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} + %reduce_sum.471 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} + ROOT %reduce_sum.472 = f32[]{:T(128)} add(%reduce_sum.470, %reduce_sum.471), metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.287.clone.clone (param_0.1346: bf16[151936,2048]) -> bf16[151936,2048,1] { - %param_0.1346 = bf16[151936,2048]{1,0:T(8,128)(2,1)} parameter(0) - ROOT %bitcast.526 = bf16[151936,2048,1]{1,0,2:T(8,128)(2,1)} bitcast(%param_0.1346), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} +%region_55.60 (reduce_sum.395: f32[], reduce_sum.396: f32[]) -> f32[] { + %reduce_sum.395 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} + %reduce_sum.396 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/jvp()/reduce_sum"} + ROOT %reduce_sum.400 = f32[]{:T(128)} add(%reduce_sum.395, %reduce_sum.396), metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.368.clone.clone (param_0.1347: f32[4,128], param_1.1542: bf16[4,128,2048], param_2.1281: bf16[2048]) -> bf16[4,128,2048] { - %param_2.1281 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(2) - %dot_general.476 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1281), dimensions={2}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1542 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %convert_element_type.1438 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%param_1.1542), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - %param_0.1347 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.2067 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_0.1347), dimensions={0,1}, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %mul.2066 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1438, %mul.2067), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/mul" stack_frame_id=0} - %convert_element_type.1437 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.2066), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/convert_element_type" stack_frame_id=0} - ROOT %dot_general.475 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.476, %convert_element_type.1437), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} +%fused_computation.475 (param_0.1547: s32[1,128], param_1.1722: bf16[128], param_2.1325: f32[128], param_3.885: f32[128], param_4.523: f32[]) -> (f32[], f32[], f32[128]) { + %param_0.1547 = s32[1,128]{1,0:T(1,128)S(1)} parameter(0) + %constant.449.clone.1 = s32[]{:T(128)} constant(0) + %broadcast.981 = s32[1,128]{1,0:T(1,128)} broadcast(%constant.449.clone.1), dimensions={}, metadata={op_name="broadcast.99"} + %ne.16 = pred[1,128]{1,0:T(4,128)(4,1)} compare(%param_0.1547, %broadcast.981), direction=NE, metadata={op_name="jit(train_step)/jvp()/ne" stack_frame_id=0} + %param_2.1325 = f32[128]{0:T(128)S(1)} parameter(2) + %log.18 = f32[128]{0:T(128)} log(%param_2.1325), metadata={op_name="jit(train_step)/jvp()/log" stack_frame_id=0} + %param_1.1722 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(1) + %reduce_max.16 = f32[128]{0:T(128)} convert(%param_1.1722), metadata={op_name="jit(train_step)/jvp()/reduce_max" stack_frame_id=0} + %add.943 = f32[128]{0:T(128)} add(%log.18, %reduce_max.16), metadata={op_name="jit(train_step)/jvp()/add" stack_frame_id=0} + %square.266 = f32[128]{0:T(128)} multiply(%add.943, %add.943), metadata={op_name="jit(train_step)/jvp()/square" stack_frame_id=0} + %constant.460.clone.25 = f32[]{:T(128)} constant(0) + %broadcast.966 = f32[128]{0:T(128)} broadcast(%constant.460.clone.25), dimensions={}, metadata={op_name="broadcast.120"} + %mul.1966 = f32[128]{0:T(128)} multiply(%square.266, %broadcast.966), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + %bitcast.617 = f32[1,128]{1,0:T(1,128)} bitcast(%mul.1966), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + %broadcast.977 = f32[1,128]{1,0:T(1,128)} broadcast(%constant.460.clone.25), dimensions={}, metadata={op_name="broadcast.120"} + %mul.1952 = f32[1,128]{1,0:T(1,128)} select(%ne.16, %bitcast.617, %broadcast.977), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + %reduce.242 = f32[]{:T(128)} reduce(%mul.1952, %constant.460.clone.25), dimensions={0,1}, to_apply=%region_69.74, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} + %param_3.885 = f32[128]{0:T(128)S(1)} parameter(3) + %neg.113.clone.1 = f32[128]{0:T(128)} negate(%param_3.885), metadata={op_name="jit(train_step)/jvp()/neg" stack_frame_id=0} + %bitcast.621.clone.1 = f32[1,128]{1,0:T(1,128)} bitcast(%neg.113.clone.1), metadata={op_name="jit(train_step)/jvp()/neg" stack_frame_id=0} + %add.933.clone.1 = f32[1,128]{1,0:T(1,128)} add(%bitcast.621.clone.1, %bitcast.617), metadata={op_name="jit(train_step)/jvp()/add" stack_frame_id=0} + %mul.1950.clone.1 = f32[1,128]{1,0:T(1,128)} select(%ne.16, %add.933.clone.1, %broadcast.977), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + %reduce.241.clone.1 = f32[]{:T(128)} reduce(%mul.1950.clone.1, %constant.460.clone.25), dimensions={0,1}, to_apply=%region_55.60, metadata={op_name="jit(train_step)/jvp()/reduce_sum" stack_frame_id=0} + %param_4.523 = f32[]{:T(128)S(6)} parameter(4) + %broadcast_in_dim.309.clone.1 = f32[1,128]{1,0:T(1,128)} broadcast(%param_4.523), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/broadcast_in_dim" stack_frame_id=0} + %mul.1948.clone.1 = f32[1,128]{1,0:T(1,128)} select(%ne.16, %broadcast_in_dim.309.clone.1, %broadcast.977), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %bitcast.615.clone.1 = f32[128]{0:T(128)S(1)} bitcast(%mul.1948.clone.1), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + ROOT %tuple.168 = (f32[]{:T(128)}, f32[]{:T(128)}, f32[128]{0:T(128)S(1)}) tuple(%reduce.242, %reduce.241.clone.1, %bitcast.615.clone.1) +} + +%fused_computation.363.clone.1.clone.clone (param_0.1513: bf16[128,151936], param_1.1698: f32[128], param_2.1289: f32[128], param_3.865: bf16[128], param_4.500: s32[128], param_5.434: f32[128]) -> bf16[128,151936] { + %param_5.434 = f32[128]{0:T(128)S(1)} parameter(5) + %mul.2132 = f32[1,128,151936]{2,1,0:T(8,128)} broadcast(%param_5.434), dimensions={1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %param_1.1698 = f32[128]{0:T(128)S(1)} parameter(1) + %mul.2131 = f32[1,128,151936]{2,1,0:T(8,128)} broadcast(%param_1.1698), dimensions={1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %param_0.1513 = bf16[128,151936]{1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.1370 = f32[128,151936]{1,0:T(8,128)} convert(%param_0.1513), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} + %bitcast.753 = f32[1,128,151936]{2,1,0:T(8,128)} bitcast(%convert_element_type.1370), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} + %param_3.865 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(3) + %sub.114 = f32[1,128,151936]{2,1,0:T(8,128)} broadcast(%param_3.865), dimensions={1}, metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %sub.113 = f32[1,128,151936]{2,1,0:T(8,128)} subtract(%bitcast.753, %sub.114), metadata={op_name="jit(train_step)/jvp()/sub" stack_frame_id=0} + %exp.67 = f32[1,128,151936]{2,1,0:T(8,128)} exponential(%sub.113), metadata={op_name="jit(train_step)/jvp()/exp" stack_frame_id=0} + %mul.2130 = f32[1,128,151936]{2,1,0:T(8,128)} multiply(%mul.2131, %exp.67), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %param_2.1289 = f32[128]{0:T(128)S(1)} parameter(2) + %div.982 = f32[1,128,151936]{2,1,0:T(8,128)} broadcast(%param_2.1289), dimensions={1}, metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} + %div.981 = f32[1,128,151936]{2,1,0:T(8,128)} divide(%mul.2130, %div.982), metadata={op_name="jit(train_step)/transpose(jvp())/div" stack_frame_id=0} + %param_4.500 = s32[128]{0:T(128)S(1)} parameter(4) + %eq.35 = s32[1,128,151936]{2,1,0:T(8,128)} broadcast(%param_4.500), dimensions={1}, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} + %iota.60 = s32[1,128,151936]{2,1,0:T(8,128)} iota(), iota_dimension=2, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} + %eq.34 = pred[1,128,151936]{2,1,0:T(8,128)(4,1)} compare(%eq.35, %iota.60), direction=EQ, metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/eq" stack_frame_id=0} + %convert_element_type.1369 = f32[1,128,151936]{2,1,0:T(8,128)} convert(%eq.34), metadata={op_name="jit(train_step)/jvp(jit(_one_hot))/convert_element_type" stack_frame_id=0} + %sub.112 = f32[1,128,151936]{2,1,0:T(8,128)} subtract(%div.981, %convert_element_type.1369), metadata={op_name="jit(train_step)/transpose(jvp())/sub" stack_frame_id=0} + %mul.2129 = f32[1,128,151936]{2,1,0:T(8,128)} multiply(%mul.2132, %sub.112), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %convert_element_type.1368 = bf16[1,128,151936]{2,1,0:T(8,128)(2,1)} convert(%mul.2129), metadata={op_name="jit(train_step)/transpose(jvp())/convert_element_type" stack_frame_id=0} + ROOT %bitcast.752 = bf16[128,151936]{1,0:T(8,128)(2,1)} bitcast(%convert_element_type.1368) +} + +%bitcast_fusion.2 (bitcast_input.2: bf16[151936,2048]) -> bf16[151936,2048] { + %bitcast_input.2 = bf16[151936,2048]{1,0:T(8,128)(2,1)} parameter(0) + ROOT %bitcast.760 = bf16[151936,2048]{1,0:T(8,128)(2,1)} bitcast(%bitcast_input.2) +} + +%region_10.13 (dot_general.204: bf16[], dot_general.205: bf16[]) -> bf16[] { + %dot_general.204 = bf16[]{:T(256)} parameter(0), metadata={op_name="jit(train_step)/transpose(jvp(...k,k->...k))/dot_general"} + %dot_general.205 = bf16[]{:T(256)} parameter(1), metadata={op_name="jit(train_step)/transpose(jvp(...k,k->...k))/dot_general"} + ROOT %add.416 = bf16[]{:T(256)} add(%dot_general.204, %dot_general.205), metadata={op_name="add.79"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%fused_computation.413 (param_0.1512: f32[128], param_1.1697: bf16[1,128,2048], param_2.1290: bf16[151936,2048], param_3.866: bf16[128,151936], param_4.501: f32[128], param_5.435: f32[128], param_6.299: bf16[128], param_7.189: s32[128], param_8.111: f32[128]) -> (bf16[2048], bf16[1,128,2048]) { + %param_3.866 = bf16[128,151936]{1,0:T(8,128)(2,1)S(1)} parameter(3) + %param_4.501 = f32[128]{0:T(128)S(1)} parameter(4) + %param_5.435 = f32[128]{0:T(128)S(1)} parameter(5) + %param_6.299 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(6) + %param_7.189 = s32[128]{0:T(128)S(1)} parameter(7) + %param_8.111 = f32[128]{0:T(128)S(1)} parameter(8) + %fusion.341.clone.1 = bf16[128,151936]{1,0:T(8,128)(2,1)} fusion(%param_3.866, %param_4.501, %param_5.435, %param_6.299, %param_7.189, /*index=5*/%param_8.111), kind=kLoop, calls=%fused_computation.363.clone.1.clone.clone + %param_2.1290 = bf16[151936,2048]{1,0:T(8,128)(2,1)} parameter(2) + %fusion.356 = bf16[151936,2048]{1,0:T(8,128)(2,1)} fusion(%param_2.1290), kind=kLoop, calls=%bitcast_fusion.2 + %convolution.113.clone.1 = bf16[128,2048]{1,0:T(8,128)(2,1)} convolution(%fusion.341.clone.1, %fusion.356), dim_labels=bf_io->bf, metadata={op_name="jit(train_step)/transpose(jvp())/dot_general" stack_frame_id=0} + %bitcast.593.clone.1 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%convolution.113.clone.1), metadata={op_name="jit(train_step)/transpose(jvp())/dot_general" stack_frame_id=0} + %convert.198 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%bitcast.593.clone.1) + %param_1.1697 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %convert_element_type.1302 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%param_1.1697), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} + %param_0.1512 = f32[128]{0:T(128)S(1)} parameter(0) + %mul.1899 = f32[1,128,2048]{2,1,0:T(8,128)} broadcast(%param_0.1512), dimensions={1}, metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + %mul.1898 = f32[1,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1302, %mul.1899), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + %convert_element_type.1301 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.1898), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} + %convert.199 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%convert_element_type.1301) + %multiply.436 = f32[1,128,2048]{2,1,0:T(8,128)} multiply(%convert.198, %convert.199), metadata={op_name="multiply.354"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","2048"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %convert.200 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} convert(%multiply.436) + %constant.463.clone.1 = bf16[]{:T(256)} constant(0) + %reduce.224 = bf16[2048]{0:T(1024)(128)(2,1)} reduce(%convert.200, %constant.463.clone.1), dimensions={0,1}, to_apply=%region_10.13, metadata={op_name="jit(train_step)/transpose(jvp(...k,k->...k))/dot_general" stack_frame_id=0} + ROOT %tuple.172 = (bf16[2048]{0:T(1024)(128)(2,1)}, bf16[1,128,2048]{2,1,0:T(8,128)(2,1)S(1)}) tuple(%reduce.224, %bitcast.593.clone.1) +} + +%region_12.15 (reduce_sum.198: f32[], reduce_sum.199: f32[]) -> f32[] { + %reduce_sum.198 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/transpose(jvp())/reduce_sum"} + %reduce_sum.199 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/transpose(jvp())/reduce_sum"} + ROOT %reduce_sum.200 = f32[]{:T(128)} add(%reduce_sum.198, %reduce_sum.199), metadata={op_name="jit(train_step)/transpose(jvp())/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%fused_computation.408 (param_0.1546: bf16[1,128,2048], param_1.1721: bf16[1,128,2048], param_2.1324: bf16[2048]) -> f32[128] { + %param_0.1546 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.1290 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%param_0.1546), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} + %param_1.1721 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %convert.196 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%param_1.1721) + %param_2.1324 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(2) + %dot_general.477 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1324), dimensions={2}, metadata={op_name="jit(train_step)/jvp(...k,k->...k)/dot_general" stack_frame_id=0} + %convert.197 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%dot_general.477) + %dot_general.463 = f32[1,128,2048]{2,1,0:T(8,128)} multiply(%convert.196, %convert.197), metadata={op_name="jit(train_step)/transpose(jvp(...k,k->...k))/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","2048"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %convert_element_type.1289 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%dot_general.463), metadata={op_name="jit(train_step)/transpose(jvp())/convert_element_type" stack_frame_id=0} + %mul.1884 = f32[1,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1290, %convert_element_type.1289), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %constant.460.clone.24 = f32[]{:T(128)} constant(0) + ROOT %reduce.223 = f32[128]{0:T(128)S(1)} reduce(%mul.1884, %constant.460.clone.24), dimensions={0,2}, to_apply=%region_12.15, metadata={op_name="jit(train_step)/transpose(jvp())/reduce_sum" stack_frame_id=0} +} + +%fused_computation.484 (param_0.1366: f32[128], param_1.1477: f32[128]) -> f32[128] { + %param_0.1366 = f32[128]{0:T(128)S(1)} parameter(0) + %bitcast.631 = f32[1,128]{1,0:T(1,128)} bitcast(%param_0.1366), metadata={op_name="jit(train_step)/transpose(jvp())/reduce_sum" stack_frame_id=0} + %param_1.1477 = f32[128]{0:T(128)S(1)} parameter(1) + %constant.469.clone.1 = f32[]{:T(128)} constant(0.00048828125) + %broadcast.971 = f32[128]{0:T(128)} broadcast(%constant.469.clone.1), dimensions={}, metadata={op_name="broadcast.418"} + %div.782 = f32[128]{0:T(128)} multiply(%param_1.1477, %broadcast.971), metadata={op_name="jit(train_step)/jvp()/div" stack_frame_id=0} + %constant.470.clone.1 = f32[]{:T(128)} constant(1e-06) + %broadcast.969 = f32[128]{0:T(128)} broadcast(%constant.470.clone.1), dimensions={}, metadata={op_name="jit(train_step)/jvp()/add" stack_frame_id=0} + %add.938 = f32[128]{0:T(128)} add(%div.782, %broadcast.969), metadata={op_name="jit(train_step)/jvp()/add" stack_frame_id=0} + %bitcast.630 = f32[1,128]{1,0:T(1,128)} bitcast(%add.938), metadata={op_name="jit(train_step)/jvp()/add" stack_frame_id=0} + %rsqrt.201 = f32[1,128]{1,0:T(1,128)} rsqrt(%bitcast.630), metadata={op_name="jit(train_step)/jvp()/rsqrt" stack_frame_id=0} + %div.780 = f32[1,128]{1,0:T(1,128)} divide(%rsqrt.201, %bitcast.630), metadata={op_name="jit(train_step)/jvp()/div" stack_frame_id=0} + %constant.472.clone.1 = f32[]{:T(128)} constant(-0.5) + %mul.1970 = f32[1,128]{1,0:T(1,128)} broadcast(%constant.472.clone.1), dimensions={}, metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + %mul.1964 = f32[1,128]{1,0:T(1,128)} multiply(%div.780, %mul.1970), metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + %mul.1963 = f32[1,128]{1,0:T(1,128)} multiply(%bitcast.631, %mul.1964), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %constant.473.clone.1 = f32[]{:T(128)} constant(0.0009765625) + %mul.1969 = f32[1,128]{1,0:T(1,128)} broadcast(%constant.473.clone.1), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %mul.1959 = f32[1,128]{1,0:T(1,128)} multiply(%mul.1963, %mul.1969), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + ROOT %bitcast.625 = f32[128]{0:T(128)S(1)} bitcast(%mul.1959), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} +} + +%fused_computation.405 (param_0.1174: bf16[1,128,2048], param_1.1238: f32[128], param_2.804: f32[128], param_3.435: bf16[1,128,2048], param_4.244: bf16[2048]) -> bf16[1,128,2048] { + %param_3.435 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) + %convert.194 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%param_3.435) + %param_4.244 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(4) + %dot_general.478 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_4.244), dimensions={2}, metadata={op_name="jit(train_step)/jvp(...k,k->...k)/dot_general" stack_frame_id=0} + %convert.195 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%dot_general.478) + %dot_general.464 = f32[1,128,2048]{2,1,0:T(8,128)} multiply(%convert.194, %convert.195), metadata={op_name="jit(train_step)/transpose(jvp(...k,k->...k))/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","2048"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %convert_element_type.1281 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%dot_general.464), metadata={op_name="jit(train_step)/transpose(jvp())/convert_element_type" stack_frame_id=0} + %param_2.804 = f32[128]{0:T(128)S(1)} parameter(2) + %mul.1888 = f32[1,128,2048]{2,1,0:T(8,128)} broadcast(%param_2.804), dimensions={1}, metadata={op_name="jit(train_step)/jvp()/mul" stack_frame_id=0} + %mul.1880 = f32[1,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1281, %mul.1888), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %param_0.1174 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.1292 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%param_0.1174), metadata={op_name="jit(train_step)/jvp()/convert_element_type" stack_frame_id=0} + %param_1.1238 = f32[128]{0:T(128)S(1)} parameter(1) + %mul.1887 = f32[1,128,2048]{2,1,0:T(8,128)} broadcast(%param_1.1238), dimensions={1}, metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %mul.1886 = f32[1,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1292, %mul.1887), metadata={op_name="jit(train_step)/transpose(jvp())/mul" stack_frame_id=0} + %add_any.193 = f32[1,128,2048]{2,1,0:T(8,128)} add(%mul.1880, %mul.1886), metadata={op_name="jit(train_step)/transpose(jvp())/add_any" stack_frame_id=0} + ROOT %convert_element_type.1279 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} convert(%add_any.193), metadata={op_name="jit(train_step)/transpose(jvp())/convert_element_type" stack_frame_id=0} +} + +%fused_computation.98.clone.1 (param_0.1731: bf16[4,512,6144], param_1.1840: s32[]) -> bf16[1,512,6144] { + %param_0.1731 = bf16[4,512,6144]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1840 = s32[]{:T(128)S(6)} parameter(1) + %constant.431.clone.87 = s32[]{:T(128)} constant(0) + ROOT %dynamic-slice.192 = bf16[1,512,6144]{2,1,0:T(8,128)(2,1)S(1)} dynamic-slice(%param_0.1731, %param_1.1840, %constant.431.clone.87, %constant.431.clone.87), dynamic_slice_sizes={1,512,6144}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/jit(dynamic_index_in_dim)/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} +} + +%convert_element_type.762.reduce_sub_computation (lhs.1: bf16[], rhs.1: bf16[]) -> bf16[] { + %lhs.1 = bf16[] parameter(0) + %rhs.1 = bf16[] parameter(1) + ROOT %add.793 = bf16[] add(%lhs.1, %rhs.1), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.439 (param_0.1363: bf16[151936,2048], param_1.1551: f32[4,128], param_2.1305: bf16[4,128,2048], param_3.913: bf16[2048]) -> (bf16[4,128], bf16[4,128,151936]) { - %param_1.1551 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %param_2.1305 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) - %param_3.913 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(3) - %fusion.270.clone.1 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} fusion(%param_1.1551, %param_2.1305, %param_3.913), kind=kLoop, calls=%fused_computation.368.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/decoder_norm/...k,k->...k/dot_general" stack_frame_id=0} - %param_0.1363 = bf16[151936,2048]{1,0:T(8,128)(2,1)} parameter(0) - %fusion.253.clone.1 = bf16[151936,2048,1]{1,0,2:T(8,128)(2,1)} fusion(%param_0.1363), kind=kLoop, calls=%fused_computation.287.clone.clone, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder._apply_embedding/token_embedder/convert_element_type" stack_frame_id=0} - %convolution.85.clone.1 = bf16[4,128,151936]{2,1,0:T(8,128)(2,1)} convolution(%fusion.270.clone.1, %fusion.253.clone.1), window={size=1}, dim_labels=0bf_oi0->0bf, metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/decoder.apply_output_head/dot_general" stack_frame_id=0} - %constant.1195 = bf16[]{:T(256)} constant(-inf) - %reduce.223 = bf16[4,128]{1,0:T(4,128)(2,1)S(1)} reduce(%convolution.85.clone.1, %constant.1195), dimensions={2}, to_apply=%region_8.11, metadata={op_name="jit(train_step)/jvp()/reduce_max" stack_frame_id=0} - ROOT %tuple.164 = (bf16[4,128]{1,0:T(4,128)(2,1)S(1)}, bf16[4,128,151936]{2,1,0:T(8,128)(2,1)}) tuple(%reduce.223, %convolution.85.clone.1) +%convert_element_type.758.reduce_sub_computation (lhs: bf16[], rhs: bf16[]) -> bf16[] { + %lhs = bf16[] parameter(0) + %rhs = bf16[] parameter(1) + ROOT %add.792 = bf16[] add(%lhs, %rhs), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} } -%fused_computation.440 (param_0.1358: f32[2048,4,8,128]) -> bf16[4,2048,8,128] { - %param_0.1358 = f32[2048,4,8,128]{3,2,1,0:T(8,128)S(1)} parameter(0) - %bitcast.529 = f32[4,2048,8,128]{3,2,0,1:T(8,128)} bitcast(%param_0.1358), metadata={op_name="jit(train_step)/jvp(TransformerLinenPure.apply)/TransformerLinenPure/decoder/transpose" stack_frame_id=0} - ROOT %convert.153 = bf16[4,2048,8,128]{3,2,0,1:T(8,128)(2,1)} convert(%bitcast.529) +%fused_computation.285.clone.1 (param_0.1701: bf16[4,2048], param_1.1822: s32[], param_2.1394: bf16[4,2048]) -> (bf16[2048], bf16[2048]) { + %param_0.1701 = bf16[4,2048]{1,0:T(4,128)(2,1)} parameter(0) + %param_1.1822 = s32[]{:T(128)S(6)} parameter(1) + %constant.431.clone.78 = s32[]{:T(128)} constant(0) + %dynamic_slice.299 = bf16[1,2048]{1,0:T(2,128)(2,1)} dynamic-slice(%param_0.1701, %param_1.1822, %constant.431.clone.78), dynamic_slice_sizes={1,2048}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/jit(dynamic_index_in_dim)/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %constant.1121 = bf16[]{:T(256)} constant(-0), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %reduce.263 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} reduce(%dynamic_slice.299, %constant.1121), dimensions={0}, to_apply=%convert_element_type.762.reduce_sub_computation, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %param_2.1394 = bf16[4,2048]{1,0:T(4,128)(2,1)} parameter(2) + %dynamic_slice.284.clone.3 = bf16[1,2048]{1,0:T(2,128)(2,1)} dynamic-slice(%param_2.1394, %param_1.1822, %constant.431.clone.78), dynamic_slice_sizes={1,2048}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/jit(dynamic_index_in_dim)/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %reduce.195.clone.3 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} reduce(%dynamic_slice.284.clone.3, %constant.1121), dimensions={0}, to_apply=%convert_element_type.758.reduce_sub_computation, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + ROOT %tuple.216 = (bf16[2048]{0:T(1024)(128)(2,1)S(1)}, bf16[2048]{0:T(1024)(128)(2,1)S(1)}) tuple(%reduce.263, %reduce.195.clone.3) } -%convert_element_type.767.reduce_sub_computation (lhs.1: bf16[], rhs.1: bf16[]) -> bf16[] { - %rhs.1 = bf16[] parameter(1) - %lhs.1 = bf16[] parameter(0) - ROOT %add.755 = bf16[] add(%lhs.1, %rhs.1), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +%fused_computation.120.clone.1 (param_0.1726: bf16[4,16,128,512], param_1.1838: s32[]) -> bf16[1,16,128,512] { + %param_0.1726 = bf16[4,16,128,512]{3,2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1838 = s32[]{:T(128)S(6)} parameter(1) + %constant.431.clone.86 = s32[]{:T(128)} constant(0) + ROOT %dynamic-slice.191 = bf16[1,16,128,512]{3,2,1,0:T(8,128)(2,1)S(1)} dynamic-slice(%param_0.1726, %param_1.1838, %constant.431.clone.86, %constant.431.clone.86, %constant.431.clone.86), dynamic_slice_sizes={1,16,128,512}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/jit(dynamic_index_in_dim)/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} } -%fused_computation.155.clone.clone (param_0.1534: bf16[4,2048], param_1.1687: s32[]) -> bf16[2048] { - %param_0.1534 = bf16[4,2048]{1,0:T(4,128)(2,1)} parameter(0) - %param_1.1687 = s32[]{:T(128)S(6)} parameter(1) - %constant.1361 = s32[]{:T(128)} constant(0) - %dynamic_slice.388 = bf16[1,2048]{1,0:T(2,128)(2,1)} dynamic-slice(%param_0.1534, %param_1.1687, %constant.1361), dynamic_slice_sizes={1,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - %constant.1362 = bf16[]{:T(256)} constant(-0), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %reduce.244 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} reduce(%dynamic_slice.388, %constant.1362), dimensions={0}, to_apply=%convert_element_type.767.reduce_sub_computation, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} +%fused_computation.121.clone.1 (param_0.1718: bf16[4,512,16,128], param_1.1833: s32[]) -> bf16[1,512,16,128] { + %param_0.1718 = bf16[4,512,16,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) + %param_1.1833 = s32[]{:T(128)S(6)} parameter(1) + %constant.431.clone.84 = s32[]{:T(128)} constant(0) + ROOT %dynamic-slice.189 = bf16[1,512,16,128]{1,3,2,0:T(8,128)(2,1)S(1)} dynamic-slice(%param_0.1718, %param_1.1833, %constant.431.clone.84, %constant.431.clone.84, %constant.431.clone.84), dynamic_slice_sizes={1,512,16,128}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/jit(dynamic_index_in_dim)/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} } %region_14.16 (reduce_sum.204: f32[], reduce_sum.205: f32[]) -> f32[] { - %reduce_sum.205 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/reduce_sum"} - %reduce_sum.204 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/reduce_sum"} - ROOT %reduce_sum.206 = f32[]{:T(128)} add(%reduce_sum.204, %reduce_sum.205), metadata={op_name="checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.58.clone.clone (param_0.1535: bf16[4,4,128,2048], param_1.1688: s32[]) -> f32[4,128] { - %param_0.1535 = bf16[4,4,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1688 = s32[]{:T(128)S(6)} parameter(1) - %constant.1363 = s32[]{:T(128)} constant(0) - %dynamic_slice.389 = bf16[1,4,128,2048]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1535, %param_1.1688, %constant.1363, %constant.1363, %constant.1363), dynamic_slice_sizes={1,4,128,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - %bitcast.633 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.389), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} - %convert_element_type.1564 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%bitcast.633), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %square.280 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1564, %convert_element_type.1564), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/square" stack_frame_id=0} - %constant.1364 = f32[]{:T(128)} constant(0) - ROOT %reduce.245 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%square.280, %constant.1364), dimensions={2}, to_apply=%region_14.16, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0} -} - -%fused_computation.179.clone.1.clone (param_0.1536: f32[4,128]) -> f32[4,128] { - %param_0.1536 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %constant.1366 = f32[]{:T(128)} constant(0.00048828125) - %closed_call.106 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1366), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} - %div.999 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.1536, %closed_call.106), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} - %constant.1365 = f32[]{:T(128)} constant(1e-06) - %closed_call.105 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1365), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} - %add.1039 = f32[4,128]{1,0:T(4,128)} add(%div.999, %closed_call.105), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} - ROOT %rsqrt.181 = f32[4,128]{1,0:T(4,128)S(1)} rsqrt(%add.1039), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/rsqrt" stack_frame_id=0} + %reduce_sum.204 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/reduce_sum"} + %reduce_sum.205 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/reduce_sum"} + ROOT %reduce_sum.206 = f32[]{:T(128)} add(%reduce_sum.204, %reduce_sum.205), metadata={op_name="checkpoint/rematted_computation/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%fused_computation.172.clone.1 (param_0.1699: bf16[4,1,128,2048], param_1.1821: s32[]) -> f32[128] { + %param_0.1699 = bf16[4,1,128,2048]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %param_1.1821 = s32[]{:T(128)S(6)} parameter(1) + %constant.431.clone.77 = s32[]{:T(128)} constant(0) + %dynamic-slice.184 = bf16[1,1,128,2048]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1699, %param_1.1821, %constant.431.clone.77, %constant.431.clone.77, %constant.431.clone.77), dynamic_slice_sizes={1,1,128,2048}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/jit(dynamic_index_in_dim)/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %bitcast.902 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic-slice.184), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/jit(dynamic_index_in_dim)/squeeze" stack_frame_id=0} + %convert_element_type.1459 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%bitcast.902), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %square.277 = f32[1,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1459, %convert_element_type.1459), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/square" stack_frame_id=0} + %constant.432.clone.19 = f32[]{:T(128)} constant(0) + ROOT %reduce.262 = f32[128]{0:T(128)S(1)} reduce(%square.277, %constant.432.clone.19), dimensions={0,2}, to_apply=%region_14.16, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/reduce_sum" stack_frame_id=0} +} + +%fused_computation.299.clone.1 (param_0.1700: f32[128]) -> f32[128] { + %param_0.1700 = f32[128]{0:T(128)S(1)} parameter(0) + %constant.433.clone.9 = f32[]{:T(128)} constant(0.00048828125) + %broadcast.1073 = f32[128]{0:T(128)} broadcast(%constant.433.clone.9), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call" stack_frame_id=0} + %div.1009 = f32[128]{0:T(128)} multiply(%param_0.1700, %broadcast.1073), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/div" stack_frame_id=0} + %constant.434.clone.17 = f32[]{:T(128)} constant(1e-06) + %broadcast.1072 = f32[128]{0:T(128)} broadcast(%constant.434.clone.17), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call" stack_frame_id=0} + %add.1045 = f32[128]{0:T(128)} add(%div.1009, %broadcast.1072), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/add" stack_frame_id=0} + %bitcast.904 = f32[1,128]{1,0:T(1,128)} bitcast(%add.1045), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/add" stack_frame_id=0} + %rsqrt.214 = f32[1,128]{1,0:T(1,128)} rsqrt(%bitcast.904), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/rsqrt" stack_frame_id=0} + ROOT %bitcast.903 = f32[128]{0:T(128)S(1)} bitcast(%rsqrt.214), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/rsqrt" stack_frame_id=0} +} + +%fused_computation.227.clone.clone.clone.1 (param_0.1720: bf16[2048], param_1.1834: f32[128], param_2.1402: bf16[4,1,128,2048], param_3.928: s32[]) -> bf16[128,2048,1] { + %param_0.1720 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(0) + %dot_general.593 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_0.1720), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/...k,k->...k/dot_general" stack_frame_id=0} + %convert.219 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%dot_general.593) + %param_2.1402 = bf16[4,1,128,2048]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %param_3.928 = s32[]{:T(128)S(6)} parameter(3) + %constant.431.clone.85 = s32[]{:T(128)} constant(0) + %dynamic-slice.190 = bf16[1,1,128,2048]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_2.1402, %param_3.928, %constant.431.clone.85, %constant.431.clone.85, %constant.431.clone.85), dynamic_slice_sizes={1,1,128,2048}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/jit(dynamic_index_in_dim)/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %bitcast.920 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic-slice.190), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/jit(dynamic_index_in_dim)/squeeze" stack_frame_id=0} + %convert_element_type.1469 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%bitcast.920), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %param_1.1834 = f32[128]{0:T(128)S(1)} parameter(1) + %mul.2295 = f32[1,128,2048]{2,1,0:T(8,128)} broadcast(%param_1.1834), dimensions={1}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %mul.2294 = f32[1,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1469, %mul.2295), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %convert_element_type.1468 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.2294), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %convert.220 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%convert_element_type.1468) + %dot_general.592 = f32[1,128,2048]{2,1,0:T(8,128)} multiply(%convert.219, %convert.220), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/...k,k->...k/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","2048"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %convert.221 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} convert(%dot_general.592) + ROOT %bitcast.919 = bf16[128,2048,1]{1,0,2:T(8,128)(2,1)} bitcast(%convert.221) +} + +%fused_computation.108.clone.clone.clone.1 (param_0.1719: bf16[1,2048,16,128]) -> bf16[2048,16,128] { + %param_0.1719 = bf16[1,2048,16,128]{1,3,2,0:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.918 = bf16[2048,16,128]{0,2,1:T(8,128)(2,1)} bitcast(%param_0.1719), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/dot_general" stack_frame_id=0} } %region_15.17 (reduce_sum.207: f32[], reduce_sum.211: f32[]) -> f32[] { - %reduce_sum.211 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/reduce_sum"} - %reduce_sum.207 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/reduce_sum"} - ROOT %reduce_sum.212 = f32[]{:T(128)} add(%reduce_sum.207, %reduce_sum.211), metadata={op_name="checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.25.clone.1.clone.clone.clone.clone (param_0.1550: bf16[4,2048,16,128], param_1.1698: s32[]) -> bf16[2048,16,128,1] { - %param_0.1550 = bf16[4,2048,16,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) - %param_1.1698 = s32[]{:T(128)S(6)} parameter(1) - %constant.1377 = s32[]{:T(128)} constant(0) - %dynamic_slice.395 = bf16[1,2048,16,128]{1,3,2,0:T(8,128)(2,1)} dynamic-slice(%param_0.1550, %param_1.1698, %constant.1377, %constant.1377, %constant.1377), dynamic_slice_sizes={1,2048,16,128}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.644 = bf16[2048,16,128,1]{0,2,1,3:T(8,128)(2,1)} bitcast(%dynamic_slice.395), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} -} - -%fused_computation.114.clone.clone.clone.clone (param_0.1551: f32[4,128], param_1.1699: bf16[4,4,128,2048], param_2.1405: s32[], param_3.982: bf16[2048]) -> bf16[4,128,2048,1] { - %param_3.982 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(3) - %dot_general.571 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_3.982), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1699 = bf16[4,4,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(1) - %param_2.1405 = s32[]{:T(128)S(6)} parameter(2) - %constant.1378 = s32[]{:T(128)} constant(0) - %dynamic_slice.396 = bf16[1,4,128,2048]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_1.1699, %param_2.1405, %constant.1378, %constant.1378, %constant.1378), dynamic_slice_sizes={1,4,128,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - %bitcast.646 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.396), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} - %convert_element_type.1575 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%bitcast.646), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_0.1551 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.2256 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_0.1551), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.2255 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1575, %mul.2256), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1574 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.2255), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %dot_general.570 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.571, %convert_element_type.1574), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - ROOT %bitcast.645 = bf16[4,128,2048,1]{2,1,3,0:T(8,128)(2,1)} bitcast(%dot_general.570), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} -} - -%fused_computation.61.clone.clone (param_0.1552: bf16[4,2048,16,128], param_1.1700: s32[], param_2.1406: f32[4,128], param_3.983: bf16[4,4,128,2048], param_4.604: bf16[2048]) -> (f32[4,128,16], bf16[4,128,16,128]) { - %param_2.1406 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %param_3.983 = bf16[4,4,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(3) - %param_1.1700 = s32[]{:T(128)S(6)} parameter(1) - %param_4.604 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(4) - %fusion.74.clone.3 = bf16[4,128,2048,1]{2,1,3,0:T(8,128)(2,1)} fusion(%param_2.1406, %param_3.983, %param_1.1700, %param_4.604), kind=kLoop, calls=%fused_computation.114.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_0.1552 = bf16[4,2048,16,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) - %fusion.49.clone.3 = bf16[2048,16,128,1]{0,2,1,3:T(8,128)(2,1)} fusion(%param_0.1552, %param_1.1700), kind=kLoop, calls=%fused_computation.25.clone.1.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %convolution.44.clone.3 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)S(1)} convolution(%fusion.74.clone.3, %fusion.49.clone.3), window={size=1x16 pad=0_0x15_15 rhs_reversal=0x1}, dim_labels=0bf1_i1o0->0b1f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} - %convert_element_type.1576 = f32[4,128,16,128]{3,1,2,0:T(8,128)} convert(%convolution.44.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %square.282 = f32[4,128,16,128]{3,1,2,0:T(8,128)} multiply(%convert_element_type.1576, %convert_element_type.1576), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/square" stack_frame_id=0} - %constant.1379 = f32[]{:T(128)} constant(0) - %reduce.247 = f32[4,128,16]{1,2,0:T(8,128)S(1)} reduce(%square.282, %constant.1379), dimensions={3}, to_apply=%region_15.17, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0} - ROOT %tuple.208 = (f32[4,128,16]{1,2,0:T(8,128)S(1)}, bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)S(1)}) tuple(%reduce.247, %convolution.44.clone.3) -} - -%fused_computation.151.clone.1.clone (param_0.1553: f32[4,128,16]) -> f32[4,128,16] { - %param_0.1553 = f32[4,128,16]{1,2,0:T(8,128)S(1)} parameter(0) - %constant.1380 = f32[]{:T(128)} constant(0.0078125) - %closed_call.108 = f32[4,128,16]{1,2,0:T(8,128)} broadcast(%constant.1380), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} - %div.1001 = f32[4,128,16]{1,2,0:T(8,128)} multiply(%param_0.1553, %closed_call.108), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} - %constant.1381 = f32[]{:T(128)} constant(1e-06) - %add.1044 = f32[4,128,16]{1,2,0:T(8,128)} broadcast(%constant.1381), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} - %add.1043 = f32[4,128,16]{1,2,0:T(8,128)} add(%div.1001, %add.1044), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} - ROOT %rsqrt.183 = f32[4,128,16]{1,2,0:T(8,128)S(1)} rsqrt(%add.1043), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/rsqrt" stack_frame_id=0} -} - -%fused_computation.182.clone.clone (param_0.1549: bf16[4,128], param_1.1697: s32[]) -> bf16[128] { - %param_0.1549 = bf16[4,128]{1,0:T(4,128)(2,1)} parameter(0) - %param_1.1697 = s32[]{:T(128)S(6)} parameter(1) - %constant.1376 = s32[]{:T(128)} constant(0) - %dynamic_slice.394 = bf16[1,128]{1,0:T(2,128)(2,1)} dynamic-slice(%param_0.1549, %param_1.1697, %constant.1376), dynamic_slice_sizes={1,128}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.643 = bf16[128]{0:T(256)(128)(2,1)S(1)} bitcast(%dynamic_slice.394), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} -} - -%fused_computation.121.clone.1.clone (param_0.1554: f32[4,128,16], param_1.1701: bf16[4,128,16,128], param_2.1407: bf16[128]) -> bf16[4,128,16,128] { - %param_2.1407 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(2) - %dot_general.573 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_2.1407), dimensions={3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1701 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) - %convert_element_type.1578 = f32[4,128,16,128]{3,1,2,0:T(8,128)} convert(%param_1.1701), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_0.1554 = f32[4,128,16]{1,2,0:T(8,128)S(1)} parameter(0) - %mul.2258 = f32[4,128,16,128]{3,1,2,0:T(8,128)} broadcast(%param_0.1554), dimensions={0,1,2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.2257 = f32[4,128,16,128]{3,1,2,0:T(8,128)} multiply(%convert_element_type.1578, %mul.2258), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1577 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} convert(%mul.2257), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %dot_general.572 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)S(1)} multiply(%dot_general.573, %convert_element_type.1577), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} -} - -%fused_computation.90.clone.clone (param_0.1555: bf16[4,128,16,128]) -> (bf16[4,128,16,64], bf16[4,128,16,64]) { - %param_0.1555 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) - %split.160 = bf16[4,128,16,64]{3,1,2,0:T(8,128)(2,1)} slice(%param_0.1555), slice={[0:4], [0:128], [0:16], [64:128]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/split" stack_frame_id=0} - %neg.129 = bf16[4,128,16,64]{3,1,2,0:T(8,128)(2,1)S(1)} negate(%split.160), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/neg" stack_frame_id=0} - %split.161 = bf16[4,128,16,64]{3,1,2,0:T(8,128)(2,1)S(1)} slice(%param_0.1555), slice={[0:4], [0:128], [0:16], [0:64]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/split" stack_frame_id=0} - ROOT %tuple.209 = (bf16[4,128,16,64]{3,1,2,0:T(8,128)(2,1)S(1)}, bf16[4,128,16,64]{3,1,2,0:T(8,128)(2,1)S(1)}) tuple(%neg.129, %split.161) -} - -%fused_computation.187.clone.clone () -> f32[64] { - %constant.1355 = f32[]{:T(128)} constant(1e+06) - %closed_call.104 = f32[64]{0:T(128)} broadcast(%constant.1355), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} - %iota.51 = s32[64]{0:T(128)} iota(), iota_dimension=0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/iota" stack_frame_id=0} - %constant.1354 = s32[]{:T(128)} constant(2) - %closed_call.103 = s32[64]{0:T(128)} broadcast(%constant.1354), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} - %mul.2242 = s32[64]{0:T(128)} multiply(%iota.51, %closed_call.103), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1562 = f32[64]{0:T(128)} convert(%mul.2242), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %constant.1356 = f32[]{:T(128)} constant(0.0078125) - %closed_call.102 = f32[64]{0:T(128)} broadcast(%constant.1356), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} - %div.995 = f32[64]{0:T(128)} multiply(%convert_element_type.1562, %closed_call.102), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} - ROOT %pow.38 = f32[64]{0:T(128)S(1)} power(%closed_call.104, %div.995), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/pow" stack_frame_id=0} -} - -%fused_computation.143.clone.clone (param_0.1529: f32[64], param_1.1683: f32[4,128]) -> (bf16[4,128,1,64], bf16[4,128,1,64]) { - %param_1.1683 = f32[4,128]{1,0:T(4,128)} parameter(1) - %div.998 = f32[4,128,1,64]{3,1,0,2:T(8,128)} broadcast(%param_1.1683), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} - %param_0.1529 = f32[64]{0:T(128)S(1)} parameter(0) - %div.997 = f32[4,128,1,64]{3,1,0,2:T(8,128)} broadcast(%param_0.1529), dimensions={3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} - %div.996 = f32[4,128,1,64]{3,1,0,2:T(8,128)} divide(%div.998, %div.997), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} - %cos.43 = f32[4,128,1,64]{3,1,0,2:T(8,128)} cosine(%div.996), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/cos" stack_frame_id=0} - %convert_element_type.1563 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} convert(%cos.43), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %sin.35.clone.3 = f32[4,128,1,64]{3,1,0,2:T(8,128)} sine(%div.996), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/sin" stack_frame_id=0} - %convert_element_type.1189.clone.3 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} convert(%sin.35.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %tuple.205 = (bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)}, bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)}) tuple(%convert_element_type.1563, %convert_element_type.1189.clone.3) -} - -%fused_computation.146.clone.1.clone (param_0.1530: bf16[4,128,1,64]) -> bf16[4,128,128] { - %param_0.1530 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} parameter(0) - %constant.1357 = bf16[]{:T(256)} constant(-inf) - %pad.69 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1530, %constant.1357), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - %pad.68 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1530, %constant.1357), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - %maximum.53 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} maximum(%pad.69, %pad.68), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - ROOT %bitcast.630 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%maximum.53), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} -} - -%fused_computation.145.clone.1.clone (param_0.1545: bf16[4,128,1,64]) -> bf16[4,128,128] { - %param_0.1545 = bf16[4,128,1,64]{3,1,0,2:T(8,128)(2,1)S(1)} parameter(0) - %constant.1374 = bf16[]{:T(256)} constant(-inf) - %pad.71 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1545, %constant.1374), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - %pad.70 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} pad(%param_0.1545, %constant.1374), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - %maximum.54 = bf16[4,128,1,128]{3,1,0,2:T(8,128)(2,1)} maximum(%pad.71, %pad.70), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - ROOT %bitcast.641 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%maximum.54), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} -} - -%fused_computation.94.clone.clone (param_0.1556: bf16[4,128,16,64], param_1.1702: bf16[4,128,16,64], param_2.1408: bf16[4,128,128], param_3.984: bf16[4,128,128], param_4.605: f32[4,128,16], param_5.499: bf16[4,128,16,128], param_6.384: bf16[128]) -> bf16[4,16,128,128] { - %param_6.384 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(6) - %dot_general.575 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_6.384), dimensions={3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_5.499 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(5) - %convert_element_type.1580 = f32[4,128,16,128]{3,1,2,0:T(8,128)} convert(%param_5.499), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_4.605 = f32[4,128,16]{1,2,0:T(8,128)S(1)} parameter(4) - %mul.2265 = f32[4,128,16,128]{3,1,2,0:T(8,128)} broadcast(%param_4.605), dimensions={0,1,2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.2264 = f32[4,128,16,128]{3,1,2,0:T(8,128)} multiply(%convert_element_type.1580, %mul.2265), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1579 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} convert(%mul.2264), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %dot_general.574 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} multiply(%dot_general.575, %convert_element_type.1579), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_3.984 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) - %mul.2263 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_3.984), dimensions={0,1,3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.2261 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} multiply(%dot_general.574, %mul.2263), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %param_1.1702 = bf16[4,128,16,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) - %constant.1382 = bf16[]{:T(256)} constant(-inf) - %pad.75 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_1.1702, %constant.1382), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - %param_0.1556 = bf16[4,128,16,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) - %pad.74 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1556, %constant.1382), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - %maximum.56 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} maximum(%pad.75, %pad.74), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - %param_2.1408 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) - %mul.2262 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_2.1408), dimensions={0,1,3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.2260 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} multiply(%maximum.56, %mul.2262), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %add.1045 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} add(%mul.2261, %mul.2260), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} - %constant.1383 = bf16[]{:T(256)} constant(0.08838) - %closed_call.109 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%constant.1383), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} - %mul.2259 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} multiply(%add.1045, %closed_call.109), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - ROOT %bitcast.647 = bf16[4,16,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} bitcast(%mul.2259), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} + %reduce_sum.207 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/reduce_sum"} + %reduce_sum.211 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/reduce_sum"} + ROOT %reduce_sum.212 = f32[]{:T(128)} add(%reduce_sum.207, %reduce_sum.211), metadata={op_name="checkpoint/rematted_computation/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%fused_computation.181.clone.1 (param_0.1721: bf16[1,2048,16,128], param_1.1835: bf16[2048], param_2.1403: f32[128], param_3.929: bf16[4,1,128,2048], param_4.550: s32[]) -> (f32[128,16], bf16[128,16,128]) { + %param_1.1835 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(1) + %param_2.1403 = f32[128]{0:T(128)S(1)} parameter(2) + %param_3.929 = bf16[4,1,128,2048]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(3) + %param_4.550 = s32[]{:T(128)S(6)} parameter(4) + %fusion.228.clone.3 = bf16[128,2048,1]{1,0,2:T(8,128)(2,1)} fusion(%param_1.1835, %param_2.1403, %param_3.929, %param_4.550), kind=kLoop, calls=%fused_computation.227.clone.clone.clone.1 + %param_0.1721 = bf16[1,2048,16,128]{1,3,2,0:T(8,128)(2,1)S(1)} parameter(0) + %fusion.114.clone.3 = bf16[2048,16,128]{0,2,1:T(8,128)(2,1)} fusion(%param_0.1721), kind=kLoop, calls=%fused_computation.108.clone.clone.clone.1, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/dot_general" stack_frame_id=0} + %convolution.69.clone.3 = bf16[128,16,128]{2,0,1:T(8,128)(2,1)S(1)} convolution(%fusion.228.clone.3, %fusion.114.clone.3), window={size=16 pad=15_15 rhs_reversal=1}, dim_labels=bf0_i0o->b0f, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/dot_general" stack_frame_id=0} + %convert_element_type.1470 = f32[128,16,128]{2,0,1:T(8,128)} convert(%convolution.69.clone.3), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %bitcast.921 = f32[1,128,16,128]{3,1,2,0:T(8,128)} bitcast(%convert_element_type.1470), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %square.279 = f32[1,128,16,128]{3,1,2,0:T(8,128)} multiply(%bitcast.921, %bitcast.921), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/square" stack_frame_id=0} + %constant.432.clone.21 = f32[]{:T(128)} constant(0) + %reduce.265 = f32[128,16]{0,1:T(8,128)S(1)} reduce(%square.279, %constant.432.clone.21), dimensions={0,3}, to_apply=%region_15.17, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/reduce_sum" stack_frame_id=0} + ROOT %tuple.223 = (f32[128,16]{0,1:T(8,128)S(1)}, bf16[128,16,128]{2,0,1:T(8,128)(2,1)S(1)}) tuple(%reduce.265, %convolution.69.clone.3) +} + +%fused_computation.279.clone.1 (param_0.1722: f32[128,16]) -> f32[128,16] { + %param_0.1722 = f32[128,16]{0,1:T(8,128)S(1)} parameter(0) + %constant.435.clone.13 = f32[]{:T(128)} constant(0.0078125) + %broadcast.1079 = f32[128,16]{0,1:T(8,128)} broadcast(%constant.435.clone.13), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call" stack_frame_id=0} + %div.1015 = f32[128,16]{0,1:T(8,128)} multiply(%param_0.1722, %broadcast.1079), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/div" stack_frame_id=0} + %constant.434.clone.19 = f32[]{:T(128)} constant(1e-06) + %broadcast.1078 = f32[128,16]{0,1:T(8,128)} broadcast(%constant.434.clone.19), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/add" stack_frame_id=0} + %add.1048 = f32[128,16]{0,1:T(8,128)} add(%div.1015, %broadcast.1078), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/add" stack_frame_id=0} + %bitcast.923 = f32[1,128,16]{1,2,0:T(8,128)} bitcast(%add.1048), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/add" stack_frame_id=0} + %rsqrt.216 = f32[1,128,16]{1,2,0:T(8,128)} rsqrt(%bitcast.923), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/rsqrt" stack_frame_id=0} + ROOT %bitcast.922 = f32[128,16]{0,1:T(8,128)S(1)} bitcast(%rsqrt.216), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/rsqrt" stack_frame_id=0} +} + +%fused_computation.305.clone.1 (param_0.1706: bf16[4,128], param_1.1826: s32[], param_2.1397: bf16[4,128]) -> (bf16[128], bf16[128]) { + %param_0.1706 = bf16[4,128]{1,0:T(4,128)(2,1)} parameter(0) + %param_1.1826 = s32[]{:T(128)S(6)} parameter(1) + %constant.431.clone.81 = s32[]{:T(128)} constant(0) + %dynamic_slice.300 = bf16[1,128]{1,0:T(2,128)(2,1)} dynamic-slice(%param_0.1706, %param_1.1826, %constant.431.clone.81), dynamic_slice_sizes={1,128}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/jit(dynamic_index_in_dim)/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %bitcast.909 = bf16[128]{0:T(256)(128)(2,1)S(1)} bitcast(%dynamic_slice.300), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %param_2.1397 = bf16[4,128]{1,0:T(4,128)(2,1)} parameter(2) + %dynamic_slice.294.clone.3 = bf16[1,128]{1,0:T(2,128)(2,1)} dynamic-slice(%param_2.1397, %param_1.1826, %constant.431.clone.81), dynamic_slice_sizes={1,128}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/jit(dynamic_index_in_dim)/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %bitcast.470.clone.3 = bf16[128]{0:T(256)(128)(2,1)S(1)} bitcast(%dynamic_slice.294.clone.3), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + ROOT %tuple.217 = (bf16[128]{0:T(256)(128)(2,1)S(1)}, bf16[128]{0:T(256)(128)(2,1)S(1)}) tuple(%bitcast.909, %bitcast.470.clone.3) +} + +%fused_computation.237.clone.2 (param_0.1723: f32[128,16], param_1.1836: bf16[128,16,128], param_2.1404: bf16[128]) -> bf16[1,128,16,128] { + %param_2.1404 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(2) + %dot_general.595 = bf16[1,128,16,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_2.1404), dimensions={3}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/...k,k->...k/dot_general" stack_frame_id=0} + %convert.269 = f32[1,128,16,128]{3,1,2,0:T(8,128)} convert(%dot_general.595) + %param_1.1836 = bf16[128,16,128]{2,0,1:T(8,128)(2,1)S(1)} parameter(1) + %convert_element_type.1472 = f32[128,16,128]{2,0,1:T(8,128)} convert(%param_1.1836), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %bitcast.924 = f32[1,128,16,128]{3,1,2,0:T(8,128)} bitcast(%convert_element_type.1472), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %param_0.1723 = f32[128,16]{0,1:T(8,128)S(1)} parameter(0) + %mul.2297 = f32[1,128,16,128]{3,1,2,0:T(8,128)} broadcast(%param_0.1723), dimensions={1,2}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %mul.2296 = f32[1,128,16,128]{3,1,2,0:T(8,128)} multiply(%bitcast.924, %mul.2297), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %convert_element_type.1471 = bf16[1,128,16,128]{3,1,2,0:T(8,128)(2,1)} convert(%mul.2296), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %convert.270 = f32[1,128,16,128]{3,1,2,0:T(8,128)} convert(%convert_element_type.1471) + %dot_general.594 = f32[1,128,16,128]{3,1,2,0:T(8,128)} multiply(%convert.269, %convert.270), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/...k,k->...k/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","16","128"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + ROOT %convert.271 = bf16[1,128,16,128]{3,1,2,0:T(8,128)(2,1)S(1)} convert(%dot_general.594) +} + +%fused_computation.198.clone.1 (param_0.1724: bf16[1,128,16,128]) -> (bf16[1,128,16,64], bf16[1,128,16,64]) { + %param_0.1724 = bf16[1,128,16,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %slice.75 = bf16[1,128,16,64]{3,1,2,0:T(8,128)(2,1)} slice(%param_0.1724), slice={[0:1], [0:128], [0:16], [64:128]}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/split" stack_frame_id=0} + %convert.272 = f32[1,128,16,64]{3,1,2,0:T(8,128)} convert(%slice.75) + %neg.127 = f32[1,128,16,64]{3,1,2,0:T(8,128)} negate(%convert.272), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/neg" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","16","64"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %convert.273 = bf16[1,128,16,64]{3,1,2,0:T(8,128)(2,1)S(1)} convert(%neg.127) + %slice.76 = bf16[1,128,16,64]{3,1,2,0:T(8,128)(2,1)S(1)} slice(%param_0.1724), slice={[0:1], [0:128], [0:16], [0:64]}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/split" stack_frame_id=0} + ROOT %tuple.224 = (bf16[1,128,16,64]{3,1,2,0:T(8,128)(2,1)S(1)}, bf16[1,128,16,64]{3,1,2,0:T(8,128)(2,1)S(1)}) tuple(%convert.273, %slice.76) +} + +%fused_computation.314.clone.1 () -> f32[64] { + %constant.436.clone.3 = f32[]{:T(128)} constant(1e+06) + %closed_call.68 = f32[64]{0:T(128)} broadcast(%constant.436.clone.3), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call" stack_frame_id=0} + %iota.65 = s32[64]{0:T(128)} iota(), iota_dimension=0, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/iota" stack_frame_id=0} + %constant.437.clone.3 = s32[]{:T(128)} constant(2) + %closed_call.67 = s32[64]{0:T(128)} broadcast(%constant.437.clone.3), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call" stack_frame_id=0} + %mul.2283 = s32[64]{0:T(128)} multiply(%iota.65, %closed_call.67), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %convert_element_type.1464 = f32[64]{0:T(128)} convert(%mul.2283), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %constant.435.clone.12 = f32[]{:T(128)} constant(0.0078125) + %closed_call.66 = f32[64]{0:T(128)} broadcast(%constant.435.clone.12), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call" stack_frame_id=0} + %div.1011 = f32[64]{0:T(128)} multiply(%convert_element_type.1464, %closed_call.66), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/div" stack_frame_id=0} + ROOT %pow.38 = f32[64]{0:T(128)S(1)} power(%closed_call.68, %div.1011), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/pow" stack_frame_id=0} +} + +%fused_computation.270.clone.1 (param_0.1712: f32[128], param_1.1830: f32[64]) -> (bf16[1,128,1,64], bf16[1,128,1,64]) { + %param_0.1712 = f32[128]{0:T(128)S(1)} parameter(0) + %div.1013 = f32[1,128,1,64]{3,1,2,0:T(8,128)} broadcast(%param_0.1712), dimensions={1}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/div" stack_frame_id=0} + %param_1.1830 = f32[64]{0:T(128)S(1)} parameter(1) + %div.1014 = f32[1,128,1,64]{3,1,2,0:T(8,128)} broadcast(%param_1.1830), dimensions={3}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/div" stack_frame_id=0} + %div.1012 = f32[1,128,1,64]{3,1,2,0:T(8,128)} divide(%div.1013, %div.1014), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/div" stack_frame_id=0} + %cos.43 = f32[1,128,1,64]{3,1,2,0:T(8,128)} cosine(%div.1012), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/cos" stack_frame_id=0} + %convert_element_type.1465 = bf16[1,128,1,64]{3,1,2,0:T(8,128)(2,1)S(1)} convert(%cos.43), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %sin.35.clone.3 = f32[1,128,1,64]{3,1,2,0:T(8,128)} sine(%div.1012), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/sin" stack_frame_id=0} + %convert_element_type.1140.clone.3 = bf16[1,128,1,64]{3,1,2,0:T(8,128)(2,1)S(1)} convert(%sin.35.clone.3), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + ROOT %tuple.219 = (bf16[1,128,1,64]{3,1,2,0:T(8,128)(2,1)S(1)}, bf16[1,128,1,64]{3,1,2,0:T(8,128)(2,1)S(1)}) tuple(%convert_element_type.1465, %convert_element_type.1140.clone.3) +} + +%fused_computation.275.clone.1 (param_0.1714: bf16[1,128,1,64]) -> (bf16[128,128], bf16[128,128]) { + %param_0.1714 = bf16[1,128,1,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %constant.1123 = bf16[]{:T(256)} constant(-inf) + %pad.77 = bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1714, %constant.1123), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/concatenate" stack_frame_id=0} + %convert.274 = f32[1,128,1,128]{3,1,2,0:T(8,128)} convert(%pad.77) + %pad.76 = bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1714, %constant.1123), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/concatenate" stack_frame_id=0} + %convert.275 = f32[1,128,1,128]{3,1,2,0:T(8,128)} convert(%pad.76) + %maximum.57 = f32[1,128,1,128]{3,1,2,0:T(8,128)} maximum(%convert.274, %convert.275), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/concatenate" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","1","128"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %constant.438.clone.8 = bf16[]{:T(256)} constant(0.08838) + %broadcast.1077 = bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%constant.438.clone.8), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/mul" stack_frame_id=0} + %convert.277 = f32[1,128,1,128]{3,1,2,0:T(8,128)} convert(%broadcast.1077) + %mul.2285 = f32[1,128,1,128]{3,1,2,0:T(8,128)} multiply(%maximum.57, %convert.277), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/mul" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","1","128"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %convert.278 = bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)} convert(%mul.2285) + %bitcast.916 = bf16[128,128]{1,0:T(8,128)(2,1)} bitcast(%convert.278), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/mul" stack_frame_id=0} + %convert.276 = bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)} convert(%maximum.57) + %bitcast.455.clone.3 = bf16[128,128]{1,0:T(8,128)(2,1)S(1)} bitcast(%convert.276), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + ROOT %tuple.221 = (bf16[128,128]{1,0:T(8,128)(2,1)}, bf16[128,128]{1,0:T(8,128)(2,1)S(1)}) tuple(%bitcast.916, %bitcast.455.clone.3) +} + +%fused_computation.274.clone.1 (param_0.1713: bf16[1,128,1,64]) -> (bf16[128,128], bf16[128,128]) { + %param_0.1713 = bf16[1,128,1,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %constant.1122 = bf16[]{:T(256)} constant(-inf) + %pad.75 = bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1713, %constant.1122), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/concatenate" stack_frame_id=0} + %convert.279 = f32[1,128,1,128]{3,1,2,0:T(8,128)} convert(%pad.75) + %pad.74 = bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1713, %constant.1122), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/concatenate" stack_frame_id=0} + %convert.280 = f32[1,128,1,128]{3,1,2,0:T(8,128)} convert(%pad.74) + %maximum.56 = f32[1,128,1,128]{3,1,2,0:T(8,128)} maximum(%convert.279, %convert.280), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/concatenate" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","1","128"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %constant.438.clone.7 = bf16[]{:T(256)} constant(0.08838) + %broadcast.1076 = bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%constant.438.clone.7), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/mul" stack_frame_id=0} + %convert.282 = f32[1,128,1,128]{3,1,2,0:T(8,128)} convert(%broadcast.1076) + %mul.2284 = f32[1,128,1,128]{3,1,2,0:T(8,128)} multiply(%maximum.56, %convert.282), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/mul" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","1","128"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %convert.283 = bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)} convert(%mul.2284) + %bitcast.915 = bf16[128,128]{1,0:T(8,128)(2,1)} bitcast(%convert.283), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/mul" stack_frame_id=0} + %convert.281 = bf16[1,128,1,128]{3,1,2,0:T(8,128)(2,1)} convert(%maximum.56) + %bitcast.456.clone.3 = bf16[128,128]{1,0:T(8,128)(2,1)S(1)} bitcast(%convert.281), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + ROOT %tuple.220 = (bf16[128,128]{1,0:T(8,128)(2,1)}, bf16[128,128]{1,0:T(8,128)(2,1)S(1)}) tuple(%bitcast.915, %bitcast.456.clone.3) +} + +%fused_computation.202.clone.1 (param_0.1725: bf16[1,128,16,64], param_1.1837: bf16[1,128,16,64], param_2.1405: bf16[128,128], param_3.930: bf16[128,128], param_4.551: f32[128,16], param_5.481: bf16[128,16,128], param_6.338: bf16[128]) -> bf16[16,128,128] { + %param_6.338 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(6) + %dot_general.597 = bf16[1,128,16,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_6.338), dimensions={3}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/...k,k->...k/dot_general" stack_frame_id=0} + %convert.284 = f32[1,128,16,128]{3,1,2,0:T(8,128)} convert(%dot_general.597) + %param_5.481 = bf16[128,16,128]{2,0,1:T(8,128)(2,1)S(1)} parameter(5) + %convert_element_type.1474 = f32[128,16,128]{2,0,1:T(8,128)} convert(%param_5.481), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %bitcast.926 = f32[1,128,16,128]{3,1,2,0:T(8,128)} bitcast(%convert_element_type.1474), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %param_4.551 = f32[128,16]{0,1:T(8,128)S(1)} parameter(4) + %mul.2304 = f32[1,128,16,128]{3,1,2,0:T(8,128)} broadcast(%param_4.551), dimensions={1,2}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %mul.2303 = f32[1,128,16,128]{3,1,2,0:T(8,128)} multiply(%bitcast.926, %mul.2304), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %convert_element_type.1473 = bf16[1,128,16,128]{3,1,2,0:T(8,128)(2,1)} convert(%mul.2303), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %convert.285 = f32[1,128,16,128]{3,1,2,0:T(8,128)} convert(%convert_element_type.1473) + %dot_general.596 = f32[1,128,16,128]{3,1,2,0:T(8,128)} multiply(%convert.284, %convert.285), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/...k,k->...k/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","16","128"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %param_3.930 = bf16[128,128]{1,0:T(8,128)(2,1)S(1)} parameter(3) + %mul.2302 = bf16[1,128,16,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_3.930), dimensions={1,3}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %convert.286 = f32[1,128,16,128]{3,1,2,0:T(8,128)} convert(%mul.2302) + %mul.2300 = f32[1,128,16,128]{3,1,2,0:T(8,128)} multiply(%dot_general.596, %convert.286), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","16","128"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %param_1.1837 = bf16[1,128,16,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) + %constant.1125 = bf16[]{:T(256)} constant(-inf) + %pad.81 = bf16[1,128,16,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_1.1837, %constant.1125), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/concatenate" stack_frame_id=0} + %convert.287 = f32[1,128,16,128]{3,1,2,0:T(8,128)} convert(%pad.81) + %param_0.1725 = bf16[1,128,16,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %pad.80 = bf16[1,128,16,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1725, %constant.1125), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/concatenate" stack_frame_id=0} + %convert.288 = f32[1,128,16,128]{3,1,2,0:T(8,128)} convert(%pad.80) + %maximum.59 = f32[1,128,16,128]{3,1,2,0:T(8,128)} maximum(%convert.287, %convert.288), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/concatenate" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","16","128"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %param_2.1405 = bf16[128,128]{1,0:T(8,128)(2,1)S(1)} parameter(2) + %mul.2301 = bf16[1,128,16,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_2.1405), dimensions={1,3}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %convert.289 = f32[1,128,16,128]{3,1,2,0:T(8,128)} convert(%mul.2301) + %mul.2299 = f32[1,128,16,128]{3,1,2,0:T(8,128)} multiply(%maximum.59, %convert.289), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","16","128"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %add.1049 = f32[1,128,16,128]{3,1,2,0:T(8,128)} add(%mul.2300, %mul.2299), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/add" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","16","128"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %constant.438.clone.9 = bf16[]{:T(256)} constant(0.08838) + %closed_call.69 = bf16[1,128,16,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%constant.438.clone.9), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call" stack_frame_id=0} + %convert.290 = f32[1,128,16,128]{3,1,2,0:T(8,128)} convert(%closed_call.69) + %mul.2298 = f32[1,128,16,128]{3,1,2,0:T(8,128)} multiply(%add.1049, %convert.290), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","16","128"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %convert.291 = bf16[1,128,16,128]{3,1,2,0:T(8,128)(2,1)} convert(%mul.2298) + ROOT %bitcast.925 = bf16[16,128,128]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%convert.291), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/shard_map/vmap(static_sparsity_splash)/jit(_splash_attention)/splash_mha_fwd_segmented_residuals/splash_mha_fwd_segmented_residuals/squeeze" stack_frame_id=0} +} + +%fused_computation.131.clone.1 (param_0.1707: bf16[4,512,8,128], param_1.1827: s32[]) -> bf16[1,512,8,128] { + %param_0.1707 = bf16[4,512,8,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) + %param_1.1827 = s32[]{:T(128)S(6)} parameter(1) + %constant.431.clone.82 = s32[]{:T(128)} constant(0) + ROOT %dynamic-slice.187 = bf16[1,512,8,128]{1,3,2,0:T(8,128)(2,1)S(1)} dynamic-slice(%param_0.1707, %param_1.1827, %constant.431.clone.82, %constant.431.clone.82, %constant.431.clone.82), dynamic_slice_sizes={1,512,8,128}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/jit(dynamic_index_in_dim)/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} +} + +%fused_computation.226.clone.clone.clone.1 (param_0.1709: bf16[2048], param_1.1828: f32[128], param_2.1398: bf16[4,1,128,2048], param_3.925: s32[]) -> bf16[128,2048,1] { + %param_0.1709 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(0) + %dot_general.587 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_0.1709), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/...k,k->...k/dot_general" stack_frame_id=0} + %convert.216 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%dot_general.587) + %param_2.1398 = bf16[4,1,128,2048]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %param_3.925 = s32[]{:T(128)S(6)} parameter(3) + %constant.431.clone.83 = s32[]{:T(128)} constant(0) + %dynamic-slice.188 = bf16[1,1,128,2048]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_2.1398, %param_3.925, %constant.431.clone.83, %constant.431.clone.83, %constant.431.clone.83), dynamic_slice_sizes={1,1,128,2048}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/jit(dynamic_index_in_dim)/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %bitcast.912 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic-slice.188), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/jit(dynamic_index_in_dim)/squeeze" stack_frame_id=0} + %convert_element_type.1463 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%bitcast.912), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %param_1.1828 = f32[128]{0:T(128)S(1)} parameter(1) + %mul.2282 = f32[1,128,2048]{2,1,0:T(8,128)} broadcast(%param_1.1828), dimensions={1}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %mul.2281 = f32[1,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1463, %mul.2282), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %convert_element_type.1462 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.2281), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %convert.217 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%convert_element_type.1462) + %dot_general.586 = f32[1,128,2048]{2,1,0:T(8,128)} multiply(%convert.216, %convert.217), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/...k,k->...k/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","2048"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %convert.218 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} convert(%dot_general.586) + ROOT %bitcast.911 = bf16[128,2048,1]{1,0,2:T(8,128)(2,1)} bitcast(%convert.218) +} + +%fused_computation.118.clone.clone.clone.clone.clone.clone.clone.1 (param_0.1708: bf16[1,2048,8,128]) -> bf16[2048,8,128] { + %param_0.1708 = bf16[1,2048,8,128]{1,3,2,0:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.910 = bf16[2048,8,128]{0,2,1:T(8,128)(2,1)} bitcast(%param_0.1708), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/dot_general" stack_frame_id=0} } %region_16.18 (reduce_sum.213: f32[], reduce_sum.214: f32[]) -> f32[] { - %reduce_sum.214 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/reduce_sum"} - %reduce_sum.213 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/reduce_sum"} - ROOT %reduce_sum.218 = f32[]{:T(128)} add(%reduce_sum.213, %reduce_sum.214), metadata={op_name="checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.69.clone.1.clone.clone.clone.clone (param_0.1541: bf16[4,2048,8,128], param_1.1692: s32[]) -> bf16[2048,8,128,1] { - %param_0.1541 = bf16[4,2048,8,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) - %param_1.1692 = s32[]{:T(128)S(6)} parameter(1) - %constant.1369 = s32[]{:T(128)} constant(0) - %dynamic_slice.392 = bf16[1,2048,8,128]{1,3,2,0:T(8,128)(2,1)} dynamic-slice(%param_0.1541, %param_1.1692, %constant.1369, %constant.1369, %constant.1369), dynamic_slice_sizes={1,2048,8,128}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.638 = bf16[2048,8,128,1]{0,2,1,3:T(8,128)(2,1)} bitcast(%dynamic_slice.392), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} -} - -%fused_computation.113.clone.clone.clone.clone (param_0.1542: f32[4,128], param_1.1693: bf16[4,4,128,2048], param_2.1401: s32[], param_3.979: bf16[2048]) -> bf16[4,128,2048,1] { - %param_3.979 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(3) - %dot_general.565 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_3.979), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1693 = bf16[4,4,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(1) - %param_2.1401 = s32[]{:T(128)S(6)} parameter(2) - %constant.1370 = s32[]{:T(128)} constant(0) - %dynamic_slice.393 = bf16[1,4,128,2048]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_1.1693, %param_2.1401, %constant.1370, %constant.1370, %constant.1370), dynamic_slice_sizes={1,4,128,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - %bitcast.640 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.393), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} - %convert_element_type.1568 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%bitcast.640), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_0.1542 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.2246 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_0.1542), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.2245 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1568, %mul.2246), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1567 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.2245), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %dot_general.564 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.565, %convert_element_type.1567), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - ROOT %bitcast.639 = bf16[4,128,2048,1]{2,1,3,0:T(8,128)(2,1)} bitcast(%dot_general.564), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} -} - -%fused_computation.84.clone.clone (param_0.1543: bf16[4,2048,8,128], param_1.1694: s32[], param_2.1402: f32[4,128], param_3.980: bf16[4,4,128,2048], param_4.602: bf16[2048]) -> (f32[4,128,8], bf16[4,128,8,128]) { - %param_2.1402 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %param_3.980 = bf16[4,4,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(3) - %param_1.1694 = s32[]{:T(128)S(6)} parameter(1) - %param_4.602 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(4) - %fusion.73.clone.3 = bf16[4,128,2048,1]{2,1,3,0:T(8,128)(2,1)} fusion(%param_2.1402, %param_3.980, %param_1.1694, %param_4.602), kind=kLoop, calls=%fused_computation.113.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_0.1543 = bf16[4,2048,8,128]{1,3,2,0:T(8,128)(2,1)} parameter(0) - %fusion.87.clone.3 = bf16[2048,8,128,1]{0,2,1,3:T(8,128)(2,1)} fusion(%param_0.1543, %param_1.1694), kind=kLoop, calls=%fused_computation.69.clone.1.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %convolution.50.clone.3 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} convolution(%fusion.73.clone.3, %fusion.87.clone.3), window={size=1x8 pad=0_0x7_7 rhs_reversal=0x1}, dim_labels=0bf1_i1o0->0b1f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} - %convert_element_type.1569 = f32[4,128,8,128]{3,1,2,0:T(8,128)} convert(%convolution.50.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %square.281 = f32[4,128,8,128]{3,1,2,0:T(8,128)} multiply(%convert_element_type.1569, %convert_element_type.1569), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/square" stack_frame_id=0} - %constant.1371 = f32[]{:T(128)} constant(0) - %reduce.246 = f32[4,128,8]{1,2,0:T(8,128)S(1)} reduce(%square.281, %constant.1371), dimensions={3}, to_apply=%region_16.18, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0} - ROOT %tuple.206 = (f32[4,128,8]{1,2,0:T(8,128)S(1)}, bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)}) tuple(%reduce.246, %convolution.50.clone.3) -} - -%fused_computation.154.clone.1.clone (param_0.1544: f32[4,128,8]) -> f32[4,128,8] { - %param_0.1544 = f32[4,128,8]{1,2,0:T(8,128)S(1)} parameter(0) - %constant.1372 = f32[]{:T(128)} constant(0.0078125) - %closed_call.107 = f32[4,128,8]{1,2,0:T(8,128)} broadcast(%constant.1372), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} - %div.1000 = f32[4,128,8]{1,2,0:T(8,128)} multiply(%param_0.1544, %closed_call.107), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} - %constant.1373 = f32[]{:T(128)} constant(1e-06) - %add.1041 = f32[4,128,8]{1,2,0:T(8,128)} broadcast(%constant.1373), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} - %add.1040 = f32[4,128,8]{1,2,0:T(8,128)} add(%div.1000, %add.1041), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} - ROOT %rsqrt.182 = f32[4,128,8]{1,2,0:T(8,128)S(1)} rsqrt(%add.1040), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/rsqrt" stack_frame_id=0} -} - -%fused_computation.184.clone.clone (param_0.1528: bf16[4,128], param_1.1682: s32[]) -> bf16[128] { - %param_0.1528 = bf16[4,128]{1,0:T(4,128)(2,1)} parameter(0) - %param_1.1682 = s32[]{:T(128)S(6)} parameter(1) - %constant.1353 = s32[]{:T(128)} constant(0) - %dynamic_slice.385 = bf16[1,128]{1,0:T(2,128)(2,1)} dynamic-slice(%param_0.1528, %param_1.1682, %constant.1353), dynamic_slice_sizes={1,128}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.629 = bf16[128]{0:T(256)(128)(2,1)S(1)} bitcast(%dynamic_slice.385), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} -} - -%fused_computation.139.clone.1.clone (param_0.1546: f32[4,128,8], param_1.1695: bf16[4,128,8,128], param_2.1403: bf16[128]) -> bf16[4,128,8,128] { - %param_2.1403 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(2) - %dot_general.567 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_2.1403), dimensions={3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1695 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) - %convert_element_type.1571 = f32[4,128,8,128]{3,1,2,0:T(8,128)} convert(%param_1.1695), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_0.1546 = f32[4,128,8]{1,2,0:T(8,128)S(1)} parameter(0) - %mul.2248 = f32[4,128,8,128]{3,1,2,0:T(8,128)} broadcast(%param_0.1546), dimensions={0,1,2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.2247 = f32[4,128,8,128]{3,1,2,0:T(8,128)} multiply(%convert_element_type.1571, %mul.2248), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1570 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} convert(%mul.2247), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %dot_general.566 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} multiply(%dot_general.567, %convert_element_type.1570), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} -} - -%fused_computation.126.clone.clone (param_0.1547: bf16[4,128,8,128]) -> (bf16[4,128,8,64], bf16[4,128,8,64]) { - %param_0.1547 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) - %split.158 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)} slice(%param_0.1547), slice={[0:4], [0:128], [0:8], [64:128]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/split" stack_frame_id=0} - %neg.128 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} negate(%split.158), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/neg" stack_frame_id=0} - %split.159 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} slice(%param_0.1547), slice={[0:4], [0:128], [0:8], [0:64]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/split" stack_frame_id=0} - ROOT %tuple.207 = (bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)}, bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)}) tuple(%neg.128, %split.159) -} - -%fused_computation.129.clone.clone (param_0.1548: bf16[4,128,8,64], param_1.1696: bf16[4,128,8,64], param_2.1404: bf16[4,128,128], param_3.981: bf16[4,128,128], param_4.603: f32[4,128,8], param_5.498: bf16[4,128,8,128], param_6.383: bf16[128]) -> bf16[4,8,128,128] { - %param_6.383 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(6) - %dot_general.569 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_6.383), dimensions={3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_5.498 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(5) - %convert_element_type.1573 = f32[4,128,8,128]{3,1,2,0:T(8,128)} convert(%param_5.498), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_4.603 = f32[4,128,8]{1,2,0:T(8,128)S(1)} parameter(4) - %mul.2254 = f32[4,128,8,128]{3,1,2,0:T(8,128)} broadcast(%param_4.603), dimensions={0,1,2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.2253 = f32[4,128,8,128]{3,1,2,0:T(8,128)} multiply(%convert_element_type.1573, %mul.2254), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1572 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} convert(%mul.2253), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %dot_general.568 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} multiply(%dot_general.569, %convert_element_type.1572), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_3.981 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) - %mul.2252 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_3.981), dimensions={0,1,3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.2250 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} multiply(%dot_general.568, %mul.2252), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %param_1.1696 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) - %constant.1375 = bf16[]{:T(256)} constant(-inf) - %pad.73 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_1.1696, %constant.1375), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - %param_0.1548 = bf16[4,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) - %pad.72 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1548, %constant.1375), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - %maximum.55 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} maximum(%pad.73, %pad.72), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/concatenate" stack_frame_id=0} - %param_2.1404 = bf16[4,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) - %mul.2251 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_2.1404), dimensions={0,1,3}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.2249 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} multiply(%maximum.55, %mul.2251), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %add.1042 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} add(%mul.2250, %mul.2249), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} - ROOT %bitcast.642 = bf16[4,8,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} bitcast(%add.1042), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} -} - -%fused_computation.169.clone.clone (param_0.1537: bf16[4,2048,8,128], param_1.1689: s32[]) -> bf16[1,2048,8,128] { - %param_0.1537 = bf16[4,2048,8,128]{3,2,0,1:T(8,128)(2,1)} parameter(0) - %param_1.1689 = s32[]{:T(128)S(6)} parameter(1) - %constant.1367 = s32[]{:T(128)} constant(0) - ROOT %dynamic_slice.390 = bf16[1,2048,8,128]{3,2,0,1:T(8,128)(2,1)S(1)} dynamic-slice(%param_0.1537, %param_1.1689, %constant.1367, %constant.1367, %constant.1367), dynamic_slice_sizes={1,2048,8,128}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} -} - -%fused_computation.70.clone.1.clone.clone.clone.clone (param_0.1538: bf16[1,2048,8,128]) -> bf16[2048,8,128,1] { - %param_0.1538 = bf16[1,2048,8,128]{3,2,0,1:T(8,128)(2,1)S(1)} parameter(0) - %copy.204 = bf16[1,2048,8,128]{3,1,2,0:T(8,128)(2,1)} copy(%param_0.1538), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0} - ROOT %bitcast.634 = bf16[2048,8,128,1]{2,0,1,3:T(8,128)(2,1)} bitcast(%copy.204), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} -} - -%fused_computation.111.clone.clone.clone.clone (param_0.1539: f32[4,128], param_1.1690: bf16[4,4,128,2048], param_2.1399: s32[], param_3.977: bf16[2048]) -> bf16[4,128,2048,1] { - %param_3.977 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(3) - %dot_general.563 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_3.977), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1690 = bf16[4,4,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(1) - %param_2.1399 = s32[]{:T(128)S(6)} parameter(2) - %constant.1368 = s32[]{:T(128)} constant(0) - %dynamic_slice.391 = bf16[1,4,128,2048]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_1.1690, %param_2.1399, %constant.1368, %constant.1368, %constant.1368), dynamic_slice_sizes={1,4,128,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - %bitcast.636 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.391), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} - %convert_element_type.1566 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%bitcast.636), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_0.1539 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.2244 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_0.1539), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.2243 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1566, %mul.2244), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1565 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.2243), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %dot_general.562 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.563, %convert_element_type.1565), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - ROOT %bitcast.635 = bf16[4,128,2048,1]{2,1,3,0:T(8,128)(2,1)} bitcast(%dot_general.562), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} -} - -%fused_computation.140.clone.clone (param_0.1540: bf16[1,2048,8,128], param_1.1691: f32[4,128], param_2.1400: bf16[4,4,128,2048], param_3.978: s32[], param_4.601: bf16[2048]) -> bf16[4,8,128,128] { - %param_1.1691 = f32[4,128]{1,0:T(4,128)S(1)} parameter(1) - %param_2.1400 = bf16[4,4,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(2) - %param_3.978 = s32[]{:T(128)S(6)} parameter(3) - %param_4.601 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(4) - %fusion.373 = bf16[4,128,2048,1]{2,1,3,0:T(8,128)(2,1)} fusion(%param_1.1691, %param_2.1400, %param_3.978, %param_4.601), kind=kLoop, calls=%fused_computation.111.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_0.1540 = bf16[1,2048,8,128]{3,2,0,1:T(8,128)(2,1)S(1)} parameter(0) - %fusion.372 = bf16[2048,8,128,1]{2,0,1,3:T(8,128)(2,1)} fusion(%param_0.1540), kind=kLoop, calls=%fused_computation.70.clone.1.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %convolution.106 = bf16[4,128,8,128]{3,1,2,0:T(8,128)(2,1)} convolution(%fusion.373, %fusion.372), window={size=1x8 pad=0_0x7_7 rhs_reversal=0x1}, dim_labels=0bf1_i1o0->0b1f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} - ROOT %bitcast.637 = bf16[4,8,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} bitcast(%convolution.106), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} -} - -%fused_computation.188.clone.clone (param_0.1578: f32[4,16,128,128]) -> (f32[4,16,128], f32[4,16,128,1]) { - %param_0.1578 = f32[4,16,128,128]{2,1,0,3:T(8,128)S(1)} parameter(0) - %slice.11 = f32[4,16,128,1]{2,1,0,3:T(8,128)S(1)} slice(%param_0.1578), slice={[0:4], [0:16], [0:128], [0:1]}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/shard_map/vmap(jit(_splash_attention))/slice" stack_frame_id=0} - %bitcast.660 = f32[4,16,128]{2,1,0:T(8,128)S(1)} bitcast(%slice.11), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/shard_map/vmap(jit(_splash_attention))/squeeze" stack_frame_id=0} - ROOT %tuple.213 = (f32[4,16,128]{2,1,0:T(8,128)S(1)}, f32[4,16,128,1]{2,1,0,3:T(8,128)S(1)}) tuple(%bitcast.660, %slice.11) + %reduce_sum.213 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/reduce_sum"} + %reduce_sum.214 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/reduce_sum"} + ROOT %reduce_sum.218 = f32[]{:T(128)} add(%reduce_sum.213, %reduce_sum.214), metadata={op_name="checkpoint/rematted_computation/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%fused_computation.232.clone.1 (param_0.1710: bf16[1,2048,8,128], param_1.1829: bf16[2048], param_2.1399: f32[128], param_3.926: bf16[4,1,128,2048], param_4.548: s32[]) -> (f32[128,8], f32[1,128,8,128]) { + %param_1.1829 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(1) + %param_2.1399 = f32[128]{0:T(128)S(1)} parameter(2) + %param_3.926 = bf16[4,1,128,2048]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(3) + %param_4.548 = s32[]{:T(128)S(6)} parameter(4) + %fusion.227.clone.3 = bf16[128,2048,1]{1,0,2:T(8,128)(2,1)} fusion(%param_1.1829, %param_2.1399, %param_3.926, %param_4.548), kind=kLoop, calls=%fused_computation.226.clone.clone.clone.1 + %param_0.1710 = bf16[1,2048,8,128]{1,3,2,0:T(8,128)(2,1)S(1)} parameter(0) + %fusion.194.clone.3 = bf16[2048,8,128]{0,2,1:T(8,128)(2,1)} fusion(%param_0.1710), kind=kLoop, calls=%fused_computation.118.clone.clone.clone.clone.clone.clone.clone.1, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/dot_general" stack_frame_id=0} + %convolution.105.clone.3 = bf16[128,8,128]{2,0,1:T(8,128)(2,1)} convolution(%fusion.227.clone.3, %fusion.194.clone.3), window={size=8 pad=7_7 rhs_reversal=1}, dim_labels=bf0_i0o->b0f, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/dot_general" stack_frame_id=0} + %convert_element_type.1090.clone.3 = f32[128,8,128]{2,0,1:T(8,128)} convert(%convolution.105.clone.3), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %bitcast.374.clone.3 = f32[1,128,8,128]{3,1,2,0:T(8,128)S(1)} bitcast(%convert_element_type.1090.clone.3), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %square.278 = f32[1,128,8,128]{3,1,2,0:T(8,128)} multiply(%bitcast.374.clone.3, %bitcast.374.clone.3), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/square" stack_frame_id=0} + %constant.432.clone.20 = f32[]{:T(128)} constant(0) + %reduce.264 = f32[128,8]{0,1:T(8,128)S(1)} reduce(%square.278, %constant.432.clone.20), dimensions={0,3}, to_apply=%region_16.18, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/reduce_sum" stack_frame_id=0} + ROOT %tuple.218 = (f32[128,8]{0,1:T(8,128)S(1)}, f32[1,128,8,128]{3,1,2,0:T(8,128)S(1)}) tuple(%reduce.264, %bitcast.374.clone.3) +} + +%fused_computation.288.clone.1 (param_0.1711: f32[128,8]) -> f32[128,8] { + %param_0.1711 = f32[128,8]{0,1:T(8,128)S(1)} parameter(0) + %constant.435.clone.11 = f32[]{:T(128)} constant(0.0078125) + %broadcast.1075 = f32[128,8]{0,1:T(8,128)} broadcast(%constant.435.clone.11), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call" stack_frame_id=0} + %div.1010 = f32[128,8]{0,1:T(8,128)} multiply(%param_0.1711, %broadcast.1075), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/div" stack_frame_id=0} + %constant.434.clone.18 = f32[]{:T(128)} constant(1e-06) + %broadcast.1074 = f32[128,8]{0,1:T(8,128)} broadcast(%constant.434.clone.18), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/add" stack_frame_id=0} + %add.1046 = f32[128,8]{0,1:T(8,128)} add(%div.1010, %broadcast.1074), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/add" stack_frame_id=0} + %bitcast.914 = f32[1,128,8]{1,2,0:T(8,128)} bitcast(%add.1046), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/add" stack_frame_id=0} + %rsqrt.215 = f32[1,128,8]{1,2,0:T(8,128)} rsqrt(%bitcast.914), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/rsqrt" stack_frame_id=0} + ROOT %bitcast.913 = f32[128,8]{0,1:T(8,128)S(1)} bitcast(%rsqrt.215), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/rsqrt" stack_frame_id=0} +} + +%fused_computation.267.clone.2 (param_0.1715: bf16[128], param_1.1831: f32[1,128,8,128], param_2.1400: f32[128,8]) -> bf16[1,128,8,128] { + %param_0.1715 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(0) + %dot_general.589 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_0.1715), dimensions={3}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/...k,k->...k/dot_general" stack_frame_id=0} + %convert.292 = f32[1,128,8,128]{3,1,2,0:T(8,128)} convert(%dot_general.589) + %param_1.1831 = f32[1,128,8,128]{3,1,2,0:T(8,128)S(1)} parameter(1) + %param_2.1400 = f32[128,8]{0,1:T(8,128)S(1)} parameter(2) + %mul.2287 = f32[1,128,8,128]{3,1,2,0:T(8,128)} broadcast(%param_2.1400), dimensions={1,2}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %mul.2286 = f32[1,128,8,128]{3,1,2,0:T(8,128)} multiply(%param_1.1831, %mul.2287), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %convert_element_type.1466 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)} convert(%mul.2286), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %convert.293 = f32[1,128,8,128]{3,1,2,0:T(8,128)} convert(%convert_element_type.1466) + %dot_general.588 = f32[1,128,8,128]{3,1,2,0:T(8,128)} multiply(%convert.292, %convert.293), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/...k,k->...k/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","8","128"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + ROOT %convert.294 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} convert(%dot_general.588) +} + +%fused_computation.246.clone.1 (param_0.1716: bf16[1,128,8,128]) -> (bf16[1,128,8,64], bf16[1,128,8,64]) { + %param_0.1716 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %slice.73 = bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)} slice(%param_0.1716), slice={[0:1], [0:128], [0:8], [64:128]}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/split" stack_frame_id=0} + %convert.295 = f32[1,128,8,64]{3,1,2,0:T(8,128)} convert(%slice.73) + %neg.126 = f32[1,128,8,64]{3,1,2,0:T(8,128)} negate(%convert.295), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/neg" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","8","64"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %convert.296 = bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} convert(%neg.126) + %slice.74 = bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} slice(%param_0.1716), slice={[0:1], [0:128], [0:8], [0:64]}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/split" stack_frame_id=0} + ROOT %tuple.222 = (bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)}, bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)}) tuple(%convert.296, %slice.74) +} + +%fused_computation.249.clone.1 (param_0.1717: bf16[1,128,8,64], param_1.1832: bf16[1,128,8,64], param_2.1401: bf16[128,128], param_3.927: bf16[128,128], param_4.549: f32[1,128,8,128], param_5.480: f32[128,8], param_6.337: bf16[128]) -> bf16[8,128,128] { + %param_6.337 = bf16[128]{0:T(256)(128)(2,1)S(1)} parameter(6) + %dot_general.591 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_6.337), dimensions={3}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/...k,k->...k/dot_general" stack_frame_id=0} + %convert.297 = f32[1,128,8,128]{3,1,2,0:T(8,128)} convert(%dot_general.591) + %param_4.549 = f32[1,128,8,128]{3,1,2,0:T(8,128)S(1)} parameter(4) + %param_5.480 = f32[128,8]{0,1:T(8,128)S(1)} parameter(5) + %mul.2293 = f32[1,128,8,128]{3,1,2,0:T(8,128)} broadcast(%param_5.480), dimensions={1,2}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %mul.2292 = f32[1,128,8,128]{3,1,2,0:T(8,128)} multiply(%param_4.549, %mul.2293), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %convert_element_type.1467 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)} convert(%mul.2292), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %convert.298 = f32[1,128,8,128]{3,1,2,0:T(8,128)} convert(%convert_element_type.1467) + %dot_general.590 = f32[1,128,8,128]{3,1,2,0:T(8,128)} multiply(%convert.297, %convert.298), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/...k,k->...k/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","8","128"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %param_3.927 = bf16[128,128]{1,0:T(8,128)(2,1)S(1)} parameter(3) + %mul.2291 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_3.927), dimensions={1,3}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %convert.299 = f32[1,128,8,128]{3,1,2,0:T(8,128)} convert(%mul.2291) + %mul.2289 = f32[1,128,8,128]{3,1,2,0:T(8,128)} multiply(%dot_general.590, %convert.299), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","8","128"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %param_1.1832 = bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(1) + %constant.1124 = bf16[]{:T(256)} constant(-inf) + %pad.79 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_1.1832, %constant.1124), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/concatenate" stack_frame_id=0} + %convert.300 = f32[1,128,8,128]{3,1,2,0:T(8,128)} convert(%pad.79) + %param_0.1717 = bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %pad.78 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_0.1717, %constant.1124), padding=0_0x0_0x0_0x64_0, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/concatenate" stack_frame_id=0} + %convert.301 = f32[1,128,8,128]{3,1,2,0:T(8,128)} convert(%pad.78) + %maximum.58 = f32[1,128,8,128]{3,1,2,0:T(8,128)} maximum(%convert.300, %convert.301), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/concatenate" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","8","128"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %param_2.1401 = bf16[128,128]{1,0:T(8,128)(2,1)S(1)} parameter(2) + %mul.2290 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)} broadcast(%param_2.1401), dimensions={1,3}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %convert.302 = f32[1,128,8,128]{3,1,2,0:T(8,128)} convert(%mul.2290) + %mul.2288 = f32[1,128,8,128]{3,1,2,0:T(8,128)} multiply(%maximum.58, %convert.302), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","8","128"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %add.1047 = f32[1,128,8,128]{3,1,2,0:T(8,128)} add(%mul.2289, %mul.2288), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/add" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","8","128"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %convert.303 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)} convert(%add.1047) + ROOT %bitcast.917 = bf16[8,128,128]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%convert.303), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/shard_map/vmap(static_sparsity_splash)/jit(_splash_attention)/splash_mha_fwd_segmented_residuals/splash_mha_fwd_segmented_residuals/squeeze" stack_frame_id=0} +} + +%fused_computation.130.clone.1 (param_0.1702: bf16[4,512,8,128], param_1.1823: s32[]) -> bf16[1,512,8,128] { + %param_0.1702 = bf16[4,512,8,128]{3,1,2,0:T(8,128)(2,1)} parameter(0) + %param_1.1823 = s32[]{:T(128)S(6)} parameter(1) + %constant.431.clone.79 = s32[]{:T(128)} constant(0) + ROOT %dynamic-slice.185 = bf16[1,512,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} dynamic-slice(%param_0.1702, %param_1.1823, %constant.431.clone.79, %constant.431.clone.79, %constant.431.clone.79), dynamic_slice_sizes={1,512,8,128}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/jit(dynamic_index_in_dim)/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} +} + +%fused_computation.225.clone.1 (param_0.1704: bf16[2048], param_1.1824: f32[128], param_2.1395: bf16[4,1,128,2048], param_3.923: s32[]) -> bf16[128,2048,1] { + %param_0.1704 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(0) + %dot_general.585 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_0.1704), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/...k,k->...k/dot_general" stack_frame_id=0} + %convert.213 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%dot_general.585) + %param_2.1395 = bf16[4,1,128,2048]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %param_3.923 = s32[]{:T(128)S(6)} parameter(3) + %constant.431.clone.80 = s32[]{:T(128)} constant(0) + %dynamic-slice.186 = bf16[1,1,128,2048]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_2.1395, %param_3.923, %constant.431.clone.80, %constant.431.clone.80, %constant.431.clone.80), dynamic_slice_sizes={1,1,128,2048}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/jit(dynamic_index_in_dim)/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %bitcast.907 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic-slice.186), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/jit(dynamic_index_in_dim)/squeeze" stack_frame_id=0} + %convert_element_type.1461 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%bitcast.907), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %param_1.1824 = f32[128]{0:T(128)S(1)} parameter(1) + %mul.2280 = f32[1,128,2048]{2,1,0:T(8,128)} broadcast(%param_1.1824), dimensions={1}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %mul.2279 = f32[1,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1461, %mul.2280), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %convert_element_type.1460 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.2279), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %convert.214 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%convert_element_type.1460) + %dot_general.584 = f32[1,128,2048]{2,1,0:T(8,128)} multiply(%convert.213, %convert.214), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/...k,k->...k/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","2048"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %convert.215 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} convert(%dot_general.584) + ROOT %bitcast.906 = bf16[128,2048,1]{1,0,2:T(8,128)(2,1)} bitcast(%convert.215) +} + +%fused_computation.116.clone.clone.clone.1 (param_0.1703: bf16[1,2048,8,128]) -> bf16[2048,8,128] { + %param_0.1703 = bf16[1,2048,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.905 = bf16[2048,8,128]{2,0,1:T(8,128)(2,1)} bitcast(%param_0.1703), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/dot_general" stack_frame_id=0} +} + +%fused_computation.190.clone.1 (param_0.1705: bf16[1,2048,8,128], param_1.1825: bf16[2048], param_2.1396: f32[128], param_3.924: bf16[4,1,128,2048], param_4.547: s32[]) -> bf16[8,128,128] { + %param_1.1825 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(1) + %param_2.1396 = f32[128]{0:T(128)S(1)} parameter(2) + %param_3.924 = bf16[4,1,128,2048]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(3) + %param_4.547 = s32[]{:T(128)S(6)} parameter(4) + %fusion.443 = bf16[128,2048,1]{1,0,2:T(8,128)(2,1)} fusion(%param_1.1825, %param_2.1396, %param_3.924, %param_4.547), kind=kLoop, calls=%fused_computation.225.clone.1 + %param_0.1705 = bf16[1,2048,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %fusion.442 = bf16[2048,8,128]{2,0,1:T(8,128)(2,1)} fusion(%param_0.1705), kind=kLoop, calls=%fused_computation.116.clone.clone.clone.1, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/dot_general" stack_frame_id=0} + %convolution.134 = bf16[128,8,128]{2,0,1:T(8,128)(2,1)} convolution(%fusion.443, %fusion.442), window={size=8 pad=7_7 rhs_reversal=1}, dim_labels=bf0_i0o->b0f, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/dot_general" stack_frame_id=0} + ROOT %bitcast.908 = bf16[8,128,128]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%convolution.134), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/shard_map/vmap(static_sparsity_splash)/jit(_splash_attention)/splash_mha_fwd_segmented_residuals/splash_mha_fwd_segmented_residuals/squeeze" stack_frame_id=0} +} + +%fused_computation.206.clone.clone.clone.1 (param_0.1728: bf16[16,128,128]) -> bf16[128,16,128] { + %param_0.1728 = bf16[16,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.928 = bf16[128,16,128]{2,0,1:T(8,128)(2,1)} bitcast(%param_0.1728) +} + +%fused_computation.106.clone.clone.clone.clone.clone.clone.clone.1 (param_0.1727: bf16[1,16,128,2048]) -> bf16[16,128,2048] { + %param_0.1727 = bf16[1,16,128,2048]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.927 = bf16[16,128,2048]{2,1,0:T(8,128)(2,1)} bitcast(%param_0.1727), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/dot_general" stack_frame_id=0} } %region_17.20 (reduce_sum.219: f32[], reduce_sum.220: f32[]) -> f32[] { - %reduce_sum.220 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/reduce_sum"} - %reduce_sum.219 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/reduce_sum"} - ROOT %reduce_sum.221 = f32[]{:T(128)} add(%reduce_sum.219, %reduce_sum.220), metadata={op_name="checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.26.clone.1.clone.clone.clone.clone.clone.clone (param_0.1557: bf16[4,16,128,2048], param_1.1703: s32[]) -> bf16[16,128,2048,1] { - %param_0.1557 = bf16[4,16,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1703 = s32[]{:T(128)S(6)} parameter(1) - %constant.1384 = s32[]{:T(128)} constant(0) - %dynamic_slice.397 = bf16[1,16,128,2048]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1557, %param_1.1703, %constant.1384, %constant.1384, %constant.1384), dynamic_slice_sizes={1,16,128,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.648 = bf16[16,128,2048,1]{2,1,0,3:T(8,128)(2,1)} bitcast(%dynamic_slice.397), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} -} - -%fused_computation.103.clone.clone.clone.clone.clone.clone (param_0.1558: bf16[4,16,128,128]) -> bf16[4,128,16,128] { - %param_0.1558 = bf16[4,16,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(0) - ROOT %bitcast.649 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} bitcast(%param_0.1558), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} -} - -%fused_computation.64.clone.clone (param_0.1559: bf16[4,16,128,2048], param_1.1704: s32[], param_2.1409: bf16[4,16,128,128], param_3.985: bf16[4,4,128,2048]) -> (f32[4,128], bf16[4,128,2048]) { - %param_3.985 = bf16[4,4,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(3) - %param_1.1704 = s32[]{:T(128)S(6)} parameter(1) - %constant.436.clone.1.clone.3 = s32[]{:T(128)} constant(0) - %dynamic_slice.242.clone.3 = bf16[1,4,128,2048]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_3.985, %param_1.1704, %constant.436.clone.1.clone.3, %constant.436.clone.1.clone.3, %constant.436.clone.1.clone.3), dynamic_slice_sizes={1,4,128,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - %bitcast.227.clone.3 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic_slice.242.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/squeeze" stack_frame_id=0} - %param_2.1409 = bf16[4,16,128,128]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(2) - %fusion.96.clone.3 = bf16[4,128,16,128]{3,1,2,0:T(8,128)(2,1)} fusion(%param_2.1409), kind=kLoop, calls=%fused_computation.103.clone.clone.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/transpose" stack_frame_id=0} - %param_0.1559 = bf16[4,16,128,2048]{3,2,1,0:T(8,128)(2,1)} parameter(0) - %fusion.95.clone.3 = bf16[16,128,2048,1]{2,1,0,3:T(8,128)(2,1)} fusion(%param_0.1559, %param_1.1704), kind=kLoop, calls=%fused_computation.26.clone.1.clone.clone.clone.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %convolution.62.clone.3 = bf16[4,128,2048,1]{2,1,3,0:T(8,128)(2,1)} convolution(%fusion.96.clone.3, %fusion.95.clone.3), window={size=1x16}, dim_labels=0b1f_1io0->0bf1, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} - %bitcast.203.clone.3 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} bitcast(%convolution.62.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} - %add.768.clone.3 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} add(%bitcast.227.clone.3, %bitcast.203.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} - %convert_element_type.1581 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%add.768.clone.3), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %square.283 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1581, %convert_element_type.1581), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/square" stack_frame_id=0} - %constant.1385 = f32[]{:T(128)} constant(0) - %reduce.248 = f32[4,128]{1,0:T(4,128)S(1)} reduce(%square.283, %constant.1385), dimensions={2}, to_apply=%region_17.20, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/reduce_sum" stack_frame_id=0} - ROOT %tuple.210 = (f32[4,128]{1,0:T(4,128)S(1)}, bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)}) tuple(%reduce.248, %add.768.clone.3) -} - -%convert_element_type.763.reduce_sub_computation (lhs: bf16[], rhs: bf16[]) -> bf16[] { - %rhs = bf16[] parameter(1) - %lhs = bf16[] parameter(0) - ROOT %add.754 = bf16[] add(%lhs, %rhs), backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} -} - -%fused_computation.156.clone.clone (param_0.1531: bf16[4,2048], param_1.1684: s32[]) -> bf16[2048] { - %param_0.1531 = bf16[4,2048]{1,0:T(4,128)(2,1)} parameter(0) - %param_1.1684 = s32[]{:T(128)S(6)} parameter(1) - %constant.1358 = s32[]{:T(128)} constant(0) - %dynamic_slice.386 = bf16[1,2048]{1,0:T(2,128)(2,1)} dynamic-slice(%param_0.1531, %param_1.1684, %constant.1358), dynamic_slice_sizes={1,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - %constant.1359 = bf16[]{:T(256)} constant(-0), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %reduce.243 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} reduce(%dynamic_slice.386, %constant.1359), dimensions={0}, to_apply=%convert_element_type.763.reduce_sub_computation, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} -} - -%fused_computation.13.clone.clone.clone (param_0.1532: bf16[4,6144,2048], param_1.1685: s32[]) -> bf16[6144,2048,1] { - %param_0.1532 = bf16[4,6144,2048]{2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1685 = s32[]{:T(128)S(6)} parameter(1) - %constant.1360 = s32[]{:T(128)} constant(0) - %dynamic_slice.387 = bf16[1,6144,2048]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1532, %param_1.1685, %constant.1360, %constant.1360), dynamic_slice_sizes={1,6144,2048}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.632 = bf16[6144,2048,1]{1,0,2:T(8,128)(2,1)} bitcast(%dynamic_slice.387), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} -} - -%bitcast_fusion.1.clone.clone (bitcast_input.4: bf16[4,128,2048]) -> bf16[4,128,2048] { - %bitcast_input.4 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) - ROOT %bitcast.631 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} bitcast(%bitcast_input.4) -} - -%fused_computation.14.clone.clone (param_0.1533: bf16[4,128,2048], param_1.1686: bf16[4,6144,2048], param_2.1398: s32[]) -> bf16[6144,4,128] { - %param_1.1686 = bf16[4,6144,2048]{2,1,0:T(8,128)(2,1)} parameter(1) - %param_2.1398 = s32[]{:T(128)S(6)} parameter(2) - %fusion.370 = bf16[6144,2048,1]{1,0,2:T(8,128)(2,1)} fusion(%param_1.1686, %param_2.1398), kind=kLoop, calls=%fused_computation.13.clone.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_0.1533 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) - %fusion.371 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} fusion(%param_0.1533), kind=kLoop, calls=%bitcast_fusion.1.clone.clone - ROOT %convolution.105 = bf16[6144,4,128]{0,2,1:T(8,128)(2,1)S(1)} convolution(%fusion.370, %fusion.371), window={size=4 pad=3_3 rhs_reversal=1}, dim_labels=bf0_0oi->b0f, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/layers/dot_general" stack_frame_id=0} -} - -%fused_computation.180.clone.1.clone (param_0.1560: f32[4,128]) -> f32[4,128] { - %param_0.1560 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %constant.1387 = f32[]{:T(128)} constant(0.00048828125) - %closed_call.111 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1387), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} - %div.1002 = f32[4,128]{1,0:T(4,128)} multiply(%param_0.1560, %closed_call.111), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/div" stack_frame_id=0} - %constant.1386 = f32[]{:T(128)} constant(1e-06) - %closed_call.110 = f32[4,128]{1,0:T(4,128)} broadcast(%constant.1386), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call" stack_frame_id=0} - %add.1046 = f32[4,128]{1,0:T(4,128)} add(%div.1002, %closed_call.110), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/add" stack_frame_id=0} - ROOT %rsqrt.184 = f32[4,128]{1,0:T(4,128)S(1)} rsqrt(%add.1046), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/rsqrt" stack_frame_id=0} -} - -%fused_computation.12.clone.1.clone.clone (param_0.1564: bf16[4,2048,6144], param_1.1708: s32[]) -> bf16[2048,6144,1] { - %param_0.1564 = bf16[4,2048,6144]{2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1708 = s32[]{:T(128)S(6)} parameter(1) - %constant.1389 = s32[]{:T(128)} constant(0) - %dynamic_slice.399 = bf16[1,2048,6144]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1564, %param_1.1708, %constant.1389, %constant.1389), dynamic_slice_sizes={1,2048,6144}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.651 = bf16[2048,6144,1]{1,0,2:T(8,128)(2,1)} bitcast(%dynamic_slice.399), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} -} - -%fused_computation.119.clone.3.clone.clone (param_0.1565: f32[4,128], param_1.1709: bf16[4,128,2048], param_2.1412: bf16[2048]) -> bf16[4,128,2048] { - %param_2.1412 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(2) - %dot_general.579 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_2.1412), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_1.1709 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %convert_element_type.1585 = f32[4,128,2048]{2,1,0:T(8,128)} convert(%param_1.1709), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - %param_0.1565 = f32[4,128]{1,0:T(4,128)S(1)} parameter(0) - %mul.2269 = f32[4,128,2048]{2,1,0:T(8,128)} broadcast(%param_0.1565), dimensions={0,1}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %mul.2268 = f32[4,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1585, %mul.2269), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/mul" stack_frame_id=0} - %convert_element_type.1584 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.2268), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %dot_general.578 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} multiply(%dot_general.579, %convert_element_type.1584), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} -} - -%fused_computation.21.clone.clone (param_0.1566: bf16[4,2048,6144], param_1.1710: s32[], param_2.1413: f32[4,128], param_3.987: bf16[4,128,2048], param_4.607: bf16[2048]) -> bf16[4,128,6144] { - %param_2.1413 = f32[4,128]{1,0:T(4,128)S(1)} parameter(2) - %param_3.987 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) - %param_4.607 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(4) - %fusion.377 = bf16[4,128,2048]{2,1,0:T(8,128)(2,1)} fusion(%param_2.1413, %param_3.987, %param_4.607), kind=kLoop, calls=%fused_computation.119.clone.3.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/...k,k->...k/dot_general" stack_frame_id=0} - %param_0.1566 = bf16[4,2048,6144]{2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1710 = s32[]{:T(128)S(6)} parameter(1) - %fusion.376 = bf16[2048,6144,1]{1,0,2:T(8,128)(2,1)} fusion(%param_0.1566, %param_1.1710), kind=kLoop, calls=%fused_computation.12.clone.1.clone.clone, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} - ROOT %convolution.108 = bf16[4,128,6144]{2,1,0:T(8,128)(2,1)S(1)} convolution(%fusion.377, %fusion.376), window={size=1}, dim_labels=0bf_io0->0bf, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/dot_general" stack_frame_id=0} -} - -%fused_computation.11.clone.1.clone.clone (param_0.1568: bf16[4,2048,6144], param_1.1712: s32[]) -> bf16[2048,6144,1] { - %param_0.1568 = bf16[4,2048,6144]{2,1,0:T(8,128)(2,1)} parameter(0) - %param_1.1712 = s32[]{:T(128)S(6)} parameter(1) - %constant.1391 = s32[]{:T(128)} constant(0) - %dynamic_slice.400 = bf16[1,2048,6144]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1568, %param_1.1712, %constant.1391, %constant.1391), dynamic_slice_sizes={1,2048,6144}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} - ROOT %bitcast.653 = bf16[2048,6144,1]{1,0,2:T(8,128)(2,1)} bitcast(%dynamic_slice.400), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/convert_element_type" stack_frame_id=0} -} - -%fused_computation.47.clone.1.clone.clone (param_0.1567: bf16[6144,4,128], param_1.1711: bf16[4,128,6144]) -> bf16[4,128,6144] { - %param_1.1711 = bf16[4,128,6144]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) - %constant.1390 = bf16[]{:T(256)} constant(1) - %jit_silu_.44 = bf16[4,128,6144]{2,1,0:T(8,128)(2,1)} broadcast(%constant.1390), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)" stack_frame_id=0} - %neg.130 = bf16[4,128,6144]{2,1,0:T(8,128)(2,1)} negate(%param_1.1711), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/neg" stack_frame_id=0} - %exp.69 = bf16[4,128,6144]{2,1,0:T(8,128)(2,1)} exponential(%neg.130), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/exp" stack_frame_id=0} - %add.1047 = bf16[4,128,6144]{2,1,0:T(8,128)(2,1)} add(%exp.69, %jit_silu_.44), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/add" stack_frame_id=0} - %div.1003 = bf16[4,128,6144]{2,1,0:T(8,128)(2,1)} divide(%jit_silu_.44, %add.1047), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/div" stack_frame_id=0} - %mul.2271 = bf16[4,128,6144]{2,1,0:T(8,128)(2,1)} multiply(%param_1.1711, %div.1003), metadata={op_name="jit(train_step)/transpose(jvp(TransformerLinenPure.apply))/TransformerLinenPure/decoder/while/body/closed_call/checkpoint/rematted_computation/layers/jit(silu)/mul" stack_frame_id=0} + %reduce_sum.219 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/reduce_sum"} + %reduce_sum.220 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/reduce_sum"} + ROOT %reduce_sum.221 = f32[]{:T(128)} add(%reduce_sum.219, %reduce_sum.220), metadata={op_name="checkpoint/rematted_computation/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%fused_computation.175.clone.1 (param_0.1729: bf16[1,16,128,2048], param_1.1839: bf16[16,128,128], param_2.1406: bf16[4,1,128,2048], param_3.931: s32[]) -> (f32[128], bf16[1,128,2048]) { + %param_2.1406 = bf16[4,1,128,2048]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %param_3.931 = s32[]{:T(128)S(6)} parameter(3) + %constant.431.clone.16.clone.3 = s32[]{:T(128)} constant(0) + %dynamic-slice.113.clone.3 = bf16[1,1,128,2048]{3,2,1,0:T(8,128)(2,1)} dynamic-slice(%param_2.1406, %param_3.931, %constant.431.clone.16.clone.3, %constant.431.clone.16.clone.3, %constant.431.clone.16.clone.3), dynamic_slice_sizes={1,1,128,2048}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/jit(dynamic_index_in_dim)/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} + %bitcast.430.clone.3 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} bitcast(%dynamic-slice.113.clone.3), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/jit(dynamic_index_in_dim)/squeeze" stack_frame_id=0} + %convert.304 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%bitcast.430.clone.3) + %param_1.1839 = bf16[16,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %fusion.209.clone.3 = bf16[128,16,128]{2,0,1:T(8,128)(2,1)} fusion(%param_1.1839), kind=kLoop, calls=%fused_computation.206.clone.clone.clone.1 + %param_0.1729 = bf16[1,16,128,2048]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %fusion.162.clone.3 = bf16[16,128,2048]{2,1,0:T(8,128)(2,1)} fusion(%param_0.1729), kind=kLoop, calls=%fused_computation.106.clone.clone.clone.clone.clone.clone.clone.1, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/dot_general" stack_frame_id=0} + %convolution.95.clone.3 = bf16[128,2048,1]{1,0,2:T(8,128)(2,1)} convolution(%fusion.209.clone.3, %fusion.162.clone.3), window={size=16}, dim_labels=b0f_0io->bf0, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/dot_general" stack_frame_id=0} + %bitcast.351.clone.3 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} bitcast(%convolution.95.clone.3), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/dot_general" stack_frame_id=0} + %convert.305 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%bitcast.351.clone.3) + %add.818.clone.3 = f32[1,128,2048]{2,1,0:T(8,128)} add(%convert.304, %convert.305), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/add" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","2048"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %convert_element_type.1475 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%add.818.clone.3), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %square.280 = f32[1,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1475, %convert_element_type.1475), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/square" stack_frame_id=0} + %constant.432.clone.22 = f32[]{:T(128)} constant(0) + %reduce.266 = f32[128]{0:T(128)S(1)} reduce(%square.280, %constant.432.clone.22), dimensions={0,2}, to_apply=%region_17.20, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/reduce_sum" stack_frame_id=0} + %convert.306 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)S(1)} convert(%add.818.clone.3) + ROOT %tuple.225 = (f32[128]{0:T(128)S(1)}, bf16[1,128,2048]{2,1,0:T(8,128)(2,1)S(1)}) tuple(%reduce.266, %convert.306) +} + +%fused_computation.297.clone.1 (param_0.1730: f32[128]) -> f32[128] { + %param_0.1730 = f32[128]{0:T(128)S(1)} parameter(0) + %constant.433.clone.10 = f32[]{:T(128)} constant(0.00048828125) + %broadcast.1081 = f32[128]{0:T(128)} broadcast(%constant.433.clone.10), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call" stack_frame_id=0} + %div.1016 = f32[128]{0:T(128)} multiply(%param_0.1730, %broadcast.1081), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/div" stack_frame_id=0} + %constant.434.clone.20 = f32[]{:T(128)} constant(1e-06) + %broadcast.1080 = f32[128]{0:T(128)} broadcast(%constant.434.clone.20), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call" stack_frame_id=0} + %add.1050 = f32[128]{0:T(128)} add(%div.1016, %broadcast.1080), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/add" stack_frame_id=0} + %bitcast.930 = f32[1,128]{1,0:T(1,128)} bitcast(%add.1050), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/add" stack_frame_id=0} + %rsqrt.217 = f32[1,128]{1,0:T(1,128)} rsqrt(%bitcast.930), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/rsqrt" stack_frame_id=0} + ROOT %bitcast.929 = f32[128]{0:T(128)S(1)} bitcast(%rsqrt.217), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/rsqrt" stack_frame_id=0} +} + +%fused_computation.217.clone.1 (param_0.1733: bf16[2048], param_1.1841: f32[128], param_2.1407: bf16[1,128,2048]) -> bf16[128,2048] { + %param_0.1733 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(0) + %dot_general.599 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_0.1733), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/...k,k->...k/dot_general" stack_frame_id=0} + %convert.222 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%dot_general.599) + %param_2.1407 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %convert_element_type.1477 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%param_2.1407), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %param_1.1841 = f32[128]{0:T(128)S(1)} parameter(1) + %mul.2306 = f32[1,128,2048]{2,1,0:T(8,128)} broadcast(%param_1.1841), dimensions={1}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %mul.2305 = f32[1,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1477, %mul.2306), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %convert_element_type.1476 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.2305), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %convert.223 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%convert_element_type.1476) + %dot_general.598 = f32[1,128,2048]{2,1,0:T(8,128)} multiply(%convert.222, %convert.223), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/...k,k->...k/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","2048"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %convert.224 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} convert(%dot_general.598) + ROOT %bitcast.932 = bf16[128,2048]{1,0:T(8,128)(2,1)} bitcast(%convert.224) +} + +%fused_computation.86.clone.clone.clone.1 (param_0.1732: bf16[1,2048,6144]) -> bf16[2048,6144] { + %param_0.1732 = bf16[1,2048,6144]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.931 = bf16[2048,6144]{1,0:T(8,128)(2,1)} bitcast(%param_0.1732), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/dot_general" stack_frame_id=0} +} + +%fused_computation.124.clone.1 (param_0.1734: bf16[1,2048,6144], param_1.1842: bf16[2048], param_2.1408: bf16[1,128,2048], param_3.932: f32[128]) -> bf16[1,128,6144] { + %param_1.1842 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(1) + %param_3.932 = f32[128]{0:T(128)S(1)} parameter(3) + %param_2.1408 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %fusion.445 = bf16[128,2048]{1,0:T(8,128)(2,1)} fusion(%param_1.1842, %param_3.932, %param_2.1408), kind=kLoop, calls=%fused_computation.217.clone.1 + %param_0.1734 = bf16[1,2048,6144]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %fusion.444 = bf16[2048,6144]{1,0:T(8,128)(2,1)} fusion(%param_0.1734), kind=kLoop, calls=%fused_computation.86.clone.clone.clone.1, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/dot_general" stack_frame_id=0} + %convolution.135 = bf16[128,6144]{1,0:T(8,128)(2,1)} convolution(%fusion.445, %fusion.444), dim_labels=bf_io->bf, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/dot_general" stack_frame_id=0} + ROOT %bitcast.933 = bf16[1,128,6144]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%convolution.135), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/dot_general" stack_frame_id=0} +} + +%fused_computation.100.clone.1 (param_0.1739: bf16[4,512,6144], param_1.1845: s32[]) -> bf16[1,512,6144] { + %param_0.1739 = bf16[4,512,6144]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1845 = s32[]{:T(128)S(6)} parameter(1) + %constant.431.clone.89 = s32[]{:T(128)} constant(0) + ROOT %dynamic-slice.194 = bf16[1,512,6144]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1739, %param_1.1845, %constant.431.clone.89, %constant.431.clone.89), dynamic_slice_sizes={1,512,6144}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/jit(dynamic_index_in_dim)/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} +} + +%fused_computation.218.clone.1 (param_0.1741: bf16[2048], param_1.1846: f32[128], param_2.1409: bf16[1,128,2048]) -> bf16[128,2048] { + %param_0.1741 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(0) + %dot_general.601 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_0.1741), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/...k,k->...k/dot_general" stack_frame_id=0} + %convert.225 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%dot_general.601) + %param_2.1409 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %convert_element_type.1479 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%param_2.1409), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %param_1.1846 = f32[128]{0:T(128)S(1)} parameter(1) + %mul.2308 = f32[1,128,2048]{2,1,0:T(8,128)} broadcast(%param_1.1846), dimensions={1}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %mul.2307 = f32[1,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1479, %mul.2308), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %convert_element_type.1478 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.2307), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %convert.226 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%convert_element_type.1478) + %dot_general.600 = f32[1,128,2048]{2,1,0:T(8,128)} multiply(%convert.225, %convert.226), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/...k,k->...k/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","2048"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %convert.227 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} convert(%dot_general.600) + ROOT %bitcast.938 = bf16[128,2048]{1,0:T(8,128)(2,1)} bitcast(%convert.227) +} + +%fused_computation.94.clone.clone.clone.1 (param_0.1740: bf16[1,2048,6144]) -> bf16[2048,6144] { + %param_0.1740 = bf16[1,2048,6144]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.937 = bf16[2048,6144]{1,0:T(8,128)(2,1)} bitcast(%param_0.1740), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/dot_general" stack_frame_id=0} +} + +%fused_computation.123.clone.1 (param_0.1742: bf16[1,2048,6144], param_1.1847: bf16[2048], param_2.1410: bf16[1,128,2048], param_3.933: f32[128]) -> bf16[1,128,6144] { + %param_1.1847 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(1) + %param_3.933 = f32[128]{0:T(128)S(1)} parameter(3) + %param_2.1410 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %fusion.449 = bf16[128,2048]{1,0:T(8,128)(2,1)} fusion(%param_1.1847, %param_3.933, %param_2.1410), kind=kLoop, calls=%fused_computation.218.clone.1 + %param_0.1742 = bf16[1,2048,6144]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %fusion.448 = bf16[2048,6144]{1,0:T(8,128)(2,1)} fusion(%param_0.1742), kind=kLoop, calls=%fused_computation.94.clone.clone.clone.1, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/dot_general" stack_frame_id=0} + %convolution.137 = bf16[128,6144]{1,0:T(8,128)(2,1)} convolution(%fusion.449, %fusion.448), dim_labels=bf_io->bf, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/dot_general" stack_frame_id=0} + ROOT %bitcast.939 = bf16[1,128,6144]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%convolution.137), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/dot_general" stack_frame_id=0} +} + +%fused_computation.144.clone.1 (param_0.1801: bf16[1,128,6144], param_1.1888: bf16[1,128,6144]) -> bf16[128,6144] { + %param_1.1888 = bf16[1,128,6144]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %convert.265 = f32[1,128,6144]{2,1,0:T(8,128)} convert(%param_1.1888) + %constant.441.clone.25 = bf16[]{:T(256)} constant(1) + %jit_silu_.45 = bf16[1,128,6144]{2,1,0:T(8,128)(2,1)} broadcast(%constant.441.clone.25), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/jit(silu)" stack_frame_id=0} + %convert.266 = f32[1,128,6144]{2,1,0:T(8,128)} convert(%jit_silu_.45) + %neg.132 = f32[1,128,6144]{2,1,0:T(8,128)} negate(%convert.265), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/jit(silu)/neg" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","6144"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %exp.78 = f32[1,128,6144]{2,1,0:T(8,128)} exponential(%neg.132), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/jit(silu)/exp" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","6144"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %add.1059 = f32[1,128,6144]{2,1,0:T(8,128)} add(%exp.78, %convert.266), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/jit(silu)/add" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","6144"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %div.1029 = f32[1,128,6144]{2,1,0:T(8,128)} divide(%convert.266, %add.1059), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/jit(silu)/div" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","6144"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %mul.2378 = f32[1,128,6144]{2,1,0:T(8,128)} multiply(%convert.265, %div.1029), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/jit(silu)/mul" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","6144"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %param_0.1801 = bf16[1,128,6144]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert.267 = f32[1,128,6144]{2,1,0:T(8,128)} convert(%param_0.1801) + %mul.2377 = f32[1,128,6144]{2,1,0:T(8,128)} multiply(%mul.2378, %convert.267), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","6144"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %convert.268 = bf16[1,128,6144]{2,1,0:T(8,128)(2,1)} convert(%mul.2377) + ROOT %bitcast.999 = bf16[128,6144]{1,0:T(8,128)(2,1)} bitcast(%convert.268) +} + +%fused_computation.212.clone.1 (param_0.1800: bf16[1,128,2048]) -> bf16[128,2048] { + %param_0.1800 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.998 = bf16[128,2048]{1,0:T(8,128)(2,1)} bitcast(%param_0.1800) +} + +%fused_computation.143.clone.1 (param_0.1802: bf16[1,128,6144], param_1.1889: bf16[1,128,2048], param_2.1441: bf16[1,128,6144]) -> bf16[6144,2048] { + %param_0.1802 = bf16[1,128,6144]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %param_2.1441 = bf16[1,128,6144]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %fusion.468 = bf16[128,6144]{1,0:T(8,128)(2,1)} fusion(%param_0.1802, %param_2.1441), kind=kLoop, calls=%fused_computation.144.clone.1 + %param_1.1889 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %fusion.469 = bf16[128,2048]{1,0:T(8,128)(2,1)} fusion(%param_1.1889), kind=kLoop, calls=%fused_computation.212.clone.1 + ROOT %convolution.147 = bf16[6144,2048]{1,0:T(8,128)(2,1)} convolution(%fusion.468, %fusion.469), dim_labels=fb_io->bf, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/dot_general" stack_frame_id=0} +} + +%fused_computation.488.clone.1 (param_0.1754: f32[1,16,128,128]) -> (f32[1,16,128,1], f32[16,128]) { + %param_0.1754 = f32[1,16,128,128]{2,1,3,0:T(8,128)S(1)} parameter(0) + %slice.77 = f32[1,16,128,1]{2,1,3,0:T(8,128)} slice(%param_0.1754), slice={[0:1], [0:16], [0:128], [0:1]}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/shard_map/vmap(static_sparsity_splash)/jit(_splash_attention)/slice" stack_frame_id=0} + %bitcast.489.clone.3 = f32[16,128]{1,0:T(8,128)S(1)} bitcast(%slice.77), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/shard_map/vmap(static_sparsity_splash)/jit(_splash_attention)/broadcast_in_dim" stack_frame_id=0} + ROOT %tuple.229 = (f32[1,16,128,1]{2,1,3,0:T(8,128)}, f32[16,128]{1,0:T(8,128)S(1)}) tuple(%slice.77, %bitcast.489.clone.3) +} + +%fused_computation.99.clone.1 (param_0.1735: bf16[4,6144,512], param_1.1843: s32[]) -> bf16[1,6144,512] { + %param_0.1735 = bf16[4,6144,512]{2,1,0:T(8,128)(2,1)} parameter(0) + %param_1.1843 = s32[]{:T(128)S(6)} parameter(1) + %constant.431.clone.88 = s32[]{:T(128)} constant(0) + ROOT %dynamic-slice.193 = bf16[1,6144,512]{2,1,0:T(8,128)(2,1)} dynamic-slice(%param_0.1735, %param_1.1843, %constant.431.clone.88, %constant.431.clone.88), dynamic_slice_sizes={1,6144,512}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/jit(dynamic_index_in_dim)/dynamic_slice" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"indices_config":{"index_known_bits":[{"zeroes":"4294967292","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"},{"zeroes":"4294967295","ones":"0","bitwidth":"32"}],"is_index_aligned":[]},"used_scoped_memory_configs":[]} +} + +%fused_computation.90.clone.clone.clone.1 (param_0.1737: bf16[1,6144,2048]) -> bf16[6144,2048] { + %param_0.1737 = bf16[1,6144,2048]{2,1,0:T(8,128)(2,1)} parameter(0) + ROOT %bitcast.935 = bf16[6144,2048]{1,0:T(8,128)(2,1)} bitcast(%param_0.1737), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/dot_general" stack_frame_id=0} +} + +%fused_computation.213.clone.1 (param_0.1736: bf16[1,128,2048]) -> bf16[128,2048] { + %param_0.1736 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} parameter(0) + ROOT %bitcast.934 = bf16[128,2048]{1,0:T(8,128)(2,1)} bitcast(%param_0.1736) +} + +%fused_computation.122.clone.1 (param_0.1738: bf16[1,6144,2048], param_1.1844: bf16[1,128,2048]) -> bf16[1,128,6144] { + %param_0.1738 = bf16[1,6144,2048]{2,1,0:T(8,128)(2,1)} parameter(0) + %fusion.446 = bf16[6144,2048]{1,0:T(8,128)(2,1)} fusion(%param_0.1738), kind=kLoop, calls=%fused_computation.90.clone.clone.clone.1, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/dot_general" stack_frame_id=0} + %param_1.1844 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} parameter(1) + %fusion.447 = bf16[128,2048]{1,0:T(8,128)(2,1)} fusion(%param_1.1844), kind=kLoop, calls=%fused_computation.213.clone.1 + %convolution.136 = bf16[6144,128]{0,1:T(8,128)(2,1)} convolution(%fusion.446, %fusion.447), dim_labels=bf_oi->bf, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/dot_general" stack_frame_id=0} + ROOT %bitcast.936 = bf16[1,128,6144]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%convolution.136), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/dot_general" stack_frame_id=0} +} + +%fused_computation.147.clone.1 (param_0.1744: bf16[1,128,6144], param_1.1848: bf16[1,128,6144]) -> bf16[128,6144] { + %param_1.1848 = bf16[1,128,6144]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %convert.228 = f32[1,128,6144]{2,1,0:T(8,128)} convert(%param_1.1848) + %constant.441.clone.21 = bf16[]{:T(256)} constant(1) + %jit_silu_.41 = bf16[1,128,6144]{2,1,0:T(8,128)(2,1)} broadcast(%constant.441.clone.21), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/jit(silu)" stack_frame_id=0} + %convert.229 = f32[1,128,6144]{2,1,0:T(8,128)} convert(%jit_silu_.41) + %neg.128 = f32[1,128,6144]{2,1,0:T(8,128)} negate(%convert.228), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/jit(silu)/neg" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","6144"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %exp.74 = f32[1,128,6144]{2,1,0:T(8,128)} exponential(%neg.128), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/jit(silu)/exp" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","6144"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %add.1051 = f32[1,128,6144]{2,1,0:T(8,128)} add(%exp.74, %convert.229), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/jit(silu)/add" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","6144"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %div.1017 = f32[1,128,6144]{2,1,0:T(8,128)} divide(%convert.229, %add.1051), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/jit(silu)/div" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","6144"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %mul.2310 = f32[1,128,6144]{2,1,0:T(8,128)} multiply(%convert.228, %div.1017), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/jit(silu)/mul" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","6144"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %param_0.1744 = bf16[1,128,6144]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert.230 = f32[1,128,6144]{2,1,0:T(8,128)} convert(%param_0.1744) + %mul.2309 = f32[1,128,6144]{2,1,0:T(8,128)} multiply(%mul.2310, %convert.230), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/mul" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","6144"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %convert.231 = bf16[1,128,6144]{2,1,0:T(8,128)(2,1)} convert(%mul.2309) + ROOT %bitcast.941 = bf16[128,6144]{1,0:T(8,128)(2,1)} bitcast(%convert.231) +} + +%fused_computation.88.clone.1 (param_0.1743: bf16[1,2048,6144]) -> bf16[2048,6144] { + %param_0.1743 = bf16[1,2048,6144]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.940 = bf16[2048,6144]{1,0:T(8,128)(2,1)} bitcast(%param_0.1743), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/dot_general" stack_frame_id=0} +} + +%fused_computation.87.clone.1 (param_0.1745: bf16[1,2048,6144], param_1.1849: bf16[1,128,6144], param_2.1411: bf16[1,128,6144]) -> bf16[128,2048] { + %param_1.1849 = bf16[1,128,6144]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %param_2.1411 = bf16[1,128,6144]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %fusion.451 = bf16[128,6144]{1,0:T(8,128)(2,1)} fusion(%param_1.1849, %param_2.1411), kind=kLoop, calls=%fused_computation.147.clone.1 + %param_0.1745 = bf16[1,2048,6144]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %fusion.450 = bf16[2048,6144]{1,0:T(8,128)(2,1)} fusion(%param_0.1745), kind=kLoop, calls=%fused_computation.88.clone.1, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/dot_general" stack_frame_id=0} + ROOT %convolution.138 = bf16[128,2048]{1,0:T(8,128)(2,1)} convolution(%fusion.451, %fusion.450), dim_labels=bf_oi->bf, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/dot_general" stack_frame_id=0} +} + +%fused_computation.142.clone.1.clone.clone.clone.1 (param_0.1747: bf16[1,128,6144], param_1.1850: bf16[1,128,6144], param_2.1412: bf16[1,128,6144]) -> bf16[128,6144] { + %param_1.1850 = bf16[1,128,6144]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %convert.232 = f32[1,128,6144]{2,1,0:T(8,128)} convert(%param_1.1850) + %param_2.1412 = bf16[1,128,6144]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %convert.233 = f32[1,128,6144]{2,1,0:T(8,128)} convert(%param_2.1412) + %mul.2315 = f32[1,128,6144]{2,1,0:T(8,128)} multiply(%convert.232, %convert.233), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/mul" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","6144"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %constant.441.clone.22 = bf16[]{:T(256)} constant(1) + %jit_silu_.42 = bf16[1,128,6144]{2,1,0:T(8,128)(2,1)} broadcast(%constant.441.clone.22), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/jit(silu)" stack_frame_id=0} + %convert.235 = f32[1,128,6144]{2,1,0:T(8,128)} convert(%jit_silu_.42) + %param_0.1747 = bf16[1,128,6144]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert.234 = f32[1,128,6144]{2,1,0:T(8,128)} convert(%param_0.1747) + %neg.129 = f32[1,128,6144]{2,1,0:T(8,128)} negate(%convert.234), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/jit(silu)/neg" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","6144"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %exp.75 = f32[1,128,6144]{2,1,0:T(8,128)} exponential(%neg.129), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/jit(silu)/exp" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","6144"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %add.1052 = f32[1,128,6144]{2,1,0:T(8,128)} add(%exp.75, %convert.235), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/jit(silu)/add" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","6144"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %div.1018 = f32[1,128,6144]{2,1,0:T(8,128)} divide(%convert.235, %add.1052), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/jit(silu)/div" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","6144"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %mul.2314 = f32[1,128,6144]{2,1,0:T(8,128)} multiply(%mul.2315, %div.1018), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/jit(silu)/mul" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","6144"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %mul.2313 = f32[1,128,6144]{2,1,0:T(8,128)} multiply(%convert.234, %mul.2315), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/jit(silu)/mul" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","6144"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %sub.118 = f32[1,128,6144]{2,1,0:T(8,128)} subtract(%convert.235, %div.1018), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/jit(silu)/sub" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","6144"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %mul.2312 = f32[1,128,6144]{2,1,0:T(8,128)} multiply(%div.1018, %sub.118), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/jit(silu)/mul" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","6144"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %mul.2311 = f32[1,128,6144]{2,1,0:T(8,128)} multiply(%mul.2313, %mul.2312), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/jit(silu)/mul" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","6144"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %add_any.208 = f32[1,128,6144]{2,1,0:T(8,128)} add(%mul.2314, %mul.2311), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/jit(silu)/add_any" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","6144"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %convert.236 = bf16[1,128,6144]{2,1,0:T(8,128)(2,1)} convert(%add_any.208) + ROOT %bitcast.943 = bf16[128,6144]{1,0:T(8,128)(2,1)} bitcast(%convert.236) +} + +%fused_computation.92.clone.clone.clone.clone.clone.clone.clone.1 (param_0.1746: bf16[1,2048,6144]) -> bf16[2048,6144] { + %param_0.1746 = bf16[1,2048,6144]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.942 = bf16[2048,6144]{1,0:T(8,128)(2,1)} bitcast(%param_0.1746), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/dot_general" stack_frame_id=0} +} + +%region_18.21 (reduce_sum.225: f32[], reduce_sum.226: f32[]) -> f32[] { + %reduce_sum.225 = f32[]{:T(128)} parameter(0), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/reduce_sum"} + %reduce_sum.226 = f32[]{:T(128)} parameter(1), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/reduce_sum"} + ROOT %reduce_sum.227 = f32[]{:T(128)} add(%reduce_sum.225, %reduce_sum.226), metadata={op_name="checkpoint/reduce_sum" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%fused_computation.177.clone.1 (param_0.1748: bf16[1,128,2048], param_1.1851: bf16[2048], param_2.1413: bf16[128,2048], param_3.934: bf16[1,2048,6144], param_4.552: bf16[1,128,6144], param_5.482: bf16[1,128,6144], param_6.339: bf16[1,128,6144]) -> (f32[128], bf16[1,128,2048]) { + %param_0.1748 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert_element_type.1481 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%param_0.1748), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %param_2.1413 = bf16[128,2048]{1,0:T(8,128)(2,1)} parameter(2) + %convert.307 = f32[128,2048]{1,0:T(8,128)} convert(%param_2.1413) + %param_4.552 = bf16[1,128,6144]{2,1,0:T(8,128)(2,1)S(1)} parameter(4) + %param_5.482 = bf16[1,128,6144]{2,1,0:T(8,128)(2,1)S(1)} parameter(5) + %param_6.339 = bf16[1,128,6144]{2,1,0:T(8,128)(2,1)S(1)} parameter(6) + %fusion.292.clone.3 = bf16[128,6144]{1,0:T(8,128)(2,1)} fusion(%param_4.552, %param_5.482, %param_6.339), kind=kLoop, calls=%fused_computation.142.clone.1.clone.clone.clone.1 + %param_3.934 = bf16[1,2048,6144]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) + %fusion.171.clone.3 = bf16[2048,6144]{1,0:T(8,128)(2,1)} fusion(%param_3.934), kind=kLoop, calls=%fused_computation.92.clone.clone.clone.clone.clone.clone.clone.1, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/dot_general" stack_frame_id=0} + %convolution.101.clone.3 = bf16[128,2048]{1,0:T(8,128)(2,1)} convolution(%fusion.292.clone.3, %fusion.171.clone.3), dim_labels=bf_oi->bf, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/dot_general" stack_frame_id=0} + %convert.308 = f32[128,2048]{1,0:T(8,128)} convert(%convolution.101.clone.3) + %add_any.151.clone.3 = f32[128,2048]{1,0:T(8,128)} add(%convert.307, %convert.308), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/add_any" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["128","2048"],"tuple_shapes":[],"layout":{"minor_to_major":["1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false]}}} + %convert.309 = bf16[128,2048]{1,0:T(8,128)(2,1)} convert(%add_any.151.clone.3) + %bitcast.359.clone.3 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)S(1)} bitcast(%convert.309), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/add_any" stack_frame_id=0} + %convert.310 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%bitcast.359.clone.3) + %param_1.1851 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(1) + %dot_general.603 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_1.1851), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/...k,k->...k/dot_general" stack_frame_id=0} + %convert.311 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%dot_general.603) + %dot_general.602 = f32[1,128,2048]{2,1,0:T(8,128)} multiply(%convert.310, %convert.311), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/...k,k->...k/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","2048"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %convert_element_type.1480 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%dot_general.602), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/convert_element_type" stack_frame_id=0} + %mul.2316 = f32[1,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1481, %convert_element_type.1480), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/mul" stack_frame_id=0} + %constant.432.clone.23 = f32[]{:T(128)} constant(0) + %reduce.267 = f32[128]{0:T(128)S(1)} reduce(%mul.2316, %constant.432.clone.23), dimensions={0,2}, to_apply=%region_18.21, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/reduce_sum" stack_frame_id=0} + ROOT %tuple.226 = (f32[128]{0:T(128)S(1)}, bf16[1,128,2048]{2,1,0:T(8,128)(2,1)S(1)}) tuple(%reduce.267, %bitcast.359.clone.3) +} + +%fused_computation.313.clone.1 (param_0.1749: f32[128], param_1.1852: f32[128]) -> f32[128] { + %param_0.1749 = f32[128]{0:T(128)S(1)} parameter(0) + %bitcast.946 = f32[1,128]{1,0:T(1,128)} bitcast(%param_0.1749), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/reduce_sum" stack_frame_id=0} + %param_1.1852 = f32[128]{0:T(128)S(1)} parameter(1) + %constant.433.clone.11 = f32[]{:T(128)} constant(0.00048828125) + %broadcast.1083 = f32[128]{0:T(128)} broadcast(%constant.433.clone.11), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call" stack_frame_id=0} + %div.1020 = f32[128]{0:T(128)} multiply(%param_1.1852, %broadcast.1083), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/div" stack_frame_id=0} + %constant.434.clone.21 = f32[]{:T(128)} constant(1e-06) + %broadcast.1082 = f32[128]{0:T(128)} broadcast(%constant.434.clone.21), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call" stack_frame_id=0} + %add.1053 = f32[128]{0:T(128)} add(%div.1020, %broadcast.1082), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/add" stack_frame_id=0} + %bitcast.945 = f32[1,128]{1,0:T(1,128)} bitcast(%add.1053), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/add" stack_frame_id=0} + %rsqrt.218 = f32[1,128]{1,0:T(1,128)} rsqrt(%bitcast.945), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/rsqrt" stack_frame_id=0} + %div.1019 = f32[1,128]{1,0:T(1,128)} divide(%rsqrt.218, %bitcast.945), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/div" stack_frame_id=0} + %constant.442.clone.9 = f32[]{:T(128)} constant(-0.5) + %closed_call.70 = f32[1,128]{1,0:T(1,128)} broadcast(%constant.442.clone.9), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call" stack_frame_id=0} + %mul.2319 = f32[1,128]{1,0:T(1,128)} multiply(%div.1019, %closed_call.70), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %mul.2318 = f32[1,128]{1,0:T(1,128)} multiply(%bitcast.946, %mul.2319), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/mul" stack_frame_id=0} + %constant.443.clone.5 = f32[]{:T(128)} constant(0.0009765625) + %mul.2320 = f32[1,128]{1,0:T(1,128)} broadcast(%constant.443.clone.5), dimensions={}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/mul" stack_frame_id=0} + %mul.2317 = f32[1,128]{1,0:T(1,128)} multiply(%mul.2318, %mul.2320), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/mul" stack_frame_id=0} + ROOT %bitcast.944 = f32[128]{0:T(128)S(1)} bitcast(%mul.2317), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/mul" stack_frame_id=0} +} + +%region_23.27 (dot_general.209: bf16[], dot_general.210: bf16[]) -> bf16[] { + %dot_general.209 = bf16[]{:T(256)} parameter(0), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/...k,k->...k/dot_general"} + %dot_general.210 = bf16[]{:T(256)} parameter(1), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/...k,k->...k/dot_general"} + ROOT %add.418 = bf16[]{:T(256)} add(%dot_general.209, %dot_general.210), metadata={op_name="add.70"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%fused_computation.230.clone.1 (param_0.1750: bf16[1,128,2048], param_1.1853: f32[128], param_2.1414: bf16[1,128,2048], param_3.935: bf16[1,128,2048], param_4.553: f32[128], param_5.483: bf16[2048]) -> (bf16[2048], bf16[1,128,2048]) { + %param_0.1750 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %convert.312 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%param_0.1750) + %param_2.1414 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %convert_element_type.1483 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%param_2.1414), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %param_1.1853 = f32[128]{0:T(128)S(1)} parameter(1) + %mul.2322 = f32[1,128,2048]{2,1,0:T(8,128)} broadcast(%param_1.1853), dimensions={1}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %mul.2321 = f32[1,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1483, %mul.2322), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %convert_element_type.1482 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} convert(%mul.2321), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/convert_element_type" stack_frame_id=0} + %convert.313 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%convert_element_type.1482) + %multiply.470 = f32[1,128,2048]{2,1,0:T(8,128)} multiply(%convert.312, %convert.313), metadata={op_name="multiply.352"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","2048"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %convert.314 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} convert(%multiply.470) + %constant.445.clone.9 = bf16[]{:T(256)} constant(0) + %reduce.268 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} reduce(%convert.314, %constant.445.clone.9), dimensions={0,1}, to_apply=%region_23.27, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/...k,k->...k/dot_general" stack_frame_id=0} + %param_3.935 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(3) + %convert.316 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%param_3.935) + %param_5.483 = bf16[2048]{0:T(1024)(128)(2,1)S(1)} parameter(5) + %dot_general.448.clone.3 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} broadcast(%param_5.483), dimensions={2}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/...k,k->...k/dot_general" stack_frame_id=0} + %convert.315 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%dot_general.448.clone.3) + %dot_general.364.clone.3 = f32[1,128,2048]{2,1,0:T(8,128)} multiply(%convert.312, %convert.315), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/...k,k->...k/dot_general" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","2048"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %convert_element_type.1115.clone.3 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%dot_general.364.clone.3), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/convert_element_type" stack_frame_id=0} + %mul.1564.clone.3 = f32[1,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1115.clone.3, %mul.2322), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/mul" stack_frame_id=0} + %param_4.553 = f32[128]{0:T(128)S(1)} parameter(4) + %mul.1583.clone.3 = f32[1,128,2048]{2,1,0:T(8,128)} broadcast(%param_4.553), dimensions={1}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/mul" stack_frame_id=0} + %mul.1563.clone.3 = f32[1,128,2048]{2,1,0:T(8,128)} multiply(%convert_element_type.1483, %mul.1583.clone.3), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/mul" stack_frame_id=0} + %add_any.171.clone.3 = f32[1,128,2048]{2,1,0:T(8,128)} add(%mul.1564.clone.3, %mul.1563.clone.3), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/add_any" stack_frame_id=0} + %convert_element_type.1113.clone.3 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)} convert(%add_any.171.clone.3), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/convert_element_type" stack_frame_id=0} + %convert.317 = f32[1,128,2048]{2,1,0:T(8,128)} convert(%convert_element_type.1113.clone.3) + %add_any.169.clone.3 = f32[1,128,2048]{2,1,0:T(8,128)} add(%convert.316, %convert.317), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/add_any" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","2048"],"tuple_shapes":[],"layout":{"minor_to_major":["2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false]}}} + %convert.318 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)S(1)} convert(%add_any.169.clone.3) + ROOT %tuple.227 = (bf16[2048]{0:T(1024)(128)(2,1)S(1)}, bf16[1,128,2048]{2,1,0:T(8,128)(2,1)S(1)}) tuple(%reduce.268, %convert.318) +} + +%fused_computation.211.clone.clone.clone.1 (param_0.1752: bf16[1,128,2048]) -> bf16[128,2048,1] { + %param_0.1752 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.948 = bf16[128,2048,1]{1,0,2:T(8,128)(2,1)} bitcast(%param_0.1752) +} + +%fused_computation.104.clone.clone.clone.1 (param_0.1751: bf16[1,16,128,2048]) -> bf16[16,128,2048] { + %param_0.1751 = bf16[1,16,128,2048]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.947 = bf16[16,128,2048]{2,1,0:T(8,128)(2,1)} bitcast(%param_0.1751), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/dot_general" stack_frame_id=0} +} + +%region_19.22 (dot_general.206: f32[], dot_general.207: f32[]) -> f32[] { + %dot_general.206 = f32[]{:T(128)} parameter(0), metadata={op_name="vmap(static_sparsity_splash)/jit(_splash_attention)/hsd,hsd->hs/dot_general"} + %dot_general.207 = f32[]{:T(128)} parameter(1), metadata={op_name="vmap(static_sparsity_splash)/jit(_splash_attention)/hsd,hsd->hs/dot_general"} + ROOT %add.417 = f32[]{:T(128)} add(%dot_general.206, %dot_general.207), metadata={op_name="add.47"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%fused_computation.183.clone.1 (param_0.1753: bf16[16,128,128], param_1.1854: bf16[1,16,128,2048], param_2.1415: bf16[1,128,2048]) -> (f32[16,128], bf16[128,16,128]) { + %param_0.1753 = bf16[16,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %bitcast.950 = bf16[1,16,128,128]{3,2,1,0:T(8,128)(2,1)} bitcast(%param_0.1753), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/shard_map/vmap(static_sparsity_splash)/jit(_splash_attention)/splash_mha_fwd_segmented_residuals/splash_mha_fwd_segmented_residuals/broadcast_in_dim" stack_frame_id=0} + %convert.193 = f32[1,16,128,128]{3,2,1,0:T(8,128)} convert(%bitcast.950), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/shard_map/convert" stack_frame_id=0} + %param_2.1415 = bf16[1,128,2048]{2,1,0:T(8,128)(2,1)S(1)} parameter(2) + %fusion.212.clone.3 = bf16[128,2048,1]{1,0,2:T(8,128)(2,1)} fusion(%param_2.1415), kind=kLoop, calls=%fused_computation.211.clone.clone.clone.1 + %param_1.1854 = bf16[1,16,128,2048]{3,2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %fusion.110.clone.3 = bf16[16,128,2048]{2,1,0:T(8,128)(2,1)} fusion(%param_1.1854), kind=kLoop, calls=%fused_computation.104.clone.clone.clone.1, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/dot_general" stack_frame_id=0} + %convolution.67.clone.3 = bf16[128,16,128]{2,0,1:T(8,128)(2,1)S(1)} convolution(%fusion.212.clone.3, %fusion.110.clone.3), window={size=16 pad=15_15 rhs_reversal=1}, dim_labels=bf0_0oi->b0f, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/dot_general" stack_frame_id=0} + %bitcast.949 = bf16[1,16,128,128]{3,2,1,0:T(8,128)(2,1)} bitcast(%convolution.67.clone.3), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/transpose" stack_frame_id=0} + %convert.192 = f32[1,16,128,128]{3,2,1,0:T(8,128)} convert(%bitcast.949), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/shard_map/convert.1" stack_frame_id=0} + %multiply.471 = f32[1,16,128,128]{3,2,1,0:T(8,128)} multiply(%convert.193, %convert.192), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/shard_map/multiply" stack_frame_id=0} + %constant.432.clone.24 = f32[]{:T(128)} constant(0) + %dot_general.604 = f32[16,128]{1,0:T(8,128)S(1)} reduce(%multiply.471, %constant.432.clone.24), dimensions={0,3}, to_apply=%region_19.22, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/shard_map/vmap(static_sparsity_splash)/jit(_splash_attention)/hsd,hsd->hs/dot_general" stack_frame_id=0} + ROOT %tuple.228 = (f32[16,128]{1,0:T(8,128)S(1)}, bf16[128,16,128]{2,0,1:T(8,128)(2,1)S(1)}) tuple(%dot_general.604, %convolution.67.clone.3) +} + +%fused_computation.252.clone.1 (param_0.1791: bf16[8,128,128]) -> bf16[128,8,128] { + %param_0.1791 = bf16[8,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.985 = bf16[128,8,128]{2,0,1:T(8,128)(2,1)} bitcast(%param_0.1791) +} + +%fused_computation.114.clone.clone.clone.1 (param_0.1790: bf16[1,2048,8,128]) -> bf16[2048,8,128] { + %param_0.1790 = bf16[1,2048,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + ROOT %bitcast.984 = bf16[2048,8,128]{2,0,1:T(8,128)(2,1)} bitcast(%param_0.1790), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/dot_general" stack_frame_id=0} +} + +%fused_computation.155.clone.1 (param_0.1792: bf16[1,2048,8,128], param_1.1882: bf16[8,128,128]) -> bf16[128,2048] { + %param_1.1882 = bf16[8,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(1) + %fusion.465 = bf16[128,8,128]{2,0,1:T(8,128)(2,1)} fusion(%param_1.1882), kind=kLoop, calls=%fused_computation.252.clone.1 + %param_0.1792 = bf16[1,2048,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %fusion.464 = bf16[2048,8,128]{2,0,1:T(8,128)(2,1)} fusion(%param_0.1792), kind=kLoop, calls=%fused_computation.114.clone.clone.clone.1, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/dot_general" stack_frame_id=0} + %convolution.145 = bf16[128,2048,1]{1,0,2:T(8,128)(2,1)} convolution(%fusion.465, %fusion.464), window={size=8}, dim_labels=b0f_o0i->bf0, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/dot_general" stack_frame_id=0} + ROOT %bitcast.986 = bf16[128,2048]{1,0:T(8,128)(2,1)S(1)} bitcast(%convolution.145), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/dot_general" stack_frame_id=0} +} + +%fused_computation.255.clone.1 (param_0.1777: bf16[8,128,128], param_1.1871: bf16[128,128]) -> bf16[1,128,8,128] { + %param_0.1777 = bf16[8,128,128]{2,1,0:T(8,128)(2,1)S(1)} parameter(0) + %bitcast.973 = bf16[1,8,128,128]{3,2,1,0:T(8,128)(2,1)} bitcast(%param_0.1777), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/shard_map/vmap(static_sparsity_splash)/jit(_splash_attention)/splash_mha_dkv_segmented_no_residuals/splash_mha_dkv_segmented_no_residuals/broadcast_in_dim" stack_frame_id=0} + %convert.319 = f32[1,8,128,128]{3,2,1,0:T(8,128)} convert(%bitcast.973) + %param_1.1871 = bf16[128,128]{1,0:T(8,128)(2,1)S(1)} parameter(1) + %broadcast.1087 = bf16[1,8,128,128]{3,2,1,0:T(8,128)(2,1)} broadcast(%param_1.1871), dimensions={2,3}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/rematted_computation/mul" stack_frame_id=0} + %convert.320 = f32[1,8,128,128]{3,2,1,0:T(8,128)} convert(%broadcast.1087) + %mul.2344 = f32[1,8,128,128]{3,2,1,0:T(8,128)} multiply(%convert.319, %convert.320), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/mul" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","8","128","128"],"tuple_shapes":[],"layout":{"minor_to_major":["3","2","1","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %convert.321 = bf16[1,8,128,128]{3,2,1,0:T(8,128)(2,1)} convert(%mul.2344) + ROOT %bitcast.972 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} bitcast(%convert.321), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/mul" stack_frame_id=0} +} + +%fused_computation.490.clone.1 (param_0.1778: bf16[1,128,8,128]) -> (bf16[1,128,8,64], bf16[1,128,8,64]) { + %param_0.1778 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(0) + %slice.79 = bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} slice(%param_0.1778), slice={[0:1], [0:128], [0:8], [64:128]}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/split" stack_frame_id=0} + %slice.52.clone.3 = bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)} slice(%param_0.1778), slice={[0:1], [0:128], [0:8], [0:64]}, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/split" stack_frame_id=0} + %convert.322 = f32[1,128,8,64]{3,1,2,0:T(8,128)} convert(%slice.52.clone.3) + %neg.100.clone.3 = f32[1,128,8,64]{3,1,2,0:T(8,128)} negate(%convert.322), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/neg" stack_frame_id=0}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"float_type_correction_info":{"original_type":"BF16","original_shape":{"element_type":"BF16","dimensions":["1","128","8","64"],"tuple_shapes":[],"layout":{"minor_to_major":["3","1","2","0"],"tiles":[{"dimensions":["8","128"]},{"dimensions":["2","1"]}],"element_size_in_bits":"0","memory_space":"0","dim_level_types":[],"index_primitive_type":"PRIMITIVE_TYPE_INVALID","pointer_primitive_type":"PRIMITIVE_TYPE_INVALID","dim_unique":[],"dim_ordered":[],"dynamic_shape_metadata_prefix_bytes":"0","tail_padding_alignment_in_elements":"1","split_configs":[]},"is_dynamic_dimension":[false,false,false,false]}}} + %convert.323 = bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} convert(%neg.100.clone.3) + ROOT %tuple.234 = (bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)}, bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)}) tuple(%slice.79, %convert.323) +} + +%region_25.29 (dot_general.213: bf16[], dot_general.214: bf16[]) -> bf16[] { + %dot_general.213 = bf16[]{:T(256)} parameter(0), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/...k,k->...k/dot_general"} + %dot_general.214 = bf16[]{:T(256)} parameter(1), metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/...k,k->...k/dot_general"} + ROOT %add.420 = bf16[]{:T(256)} add(%dot_general.213, %dot_general.214), metadata={op_name="add.72"}, backend_config={"flag_configs":[],"scoped_memory_configs":[],"used_scoped_memory_configs":[],"aliasing_operands":{"lists":[]}} +} + +%fused_computation.256.clone.1 (param_0.1779: f32[1,128,8,128], param_1.1872: f32[128,8], param_2.1428: bf16[128,128], param_3.943: bf16[8,128,128], param_4.559: bf16[1,128,8,64], param_5.486: bf16[1,128,8,64]) -> (bf16[128], bf16[1,128,8,128]) { + %param_5.486 = bf16[1,128,8,64]{3,1,2,0:T(8,128)(2,1)S(1)} parameter(5) + %constant.1098.clone.3 = bf16[]{:T(256)} constant(-inf) + %pad.61.clone.3 = bf16[1,128,8,128]{3,1,2,0:T(8,128)(2,1)} pad(%param_5.486, %constant.1098.clone.3), padding=0_0x0_0x0_0x0_64, metadata={op_name="jit(train_step)/transpose(jvp())/while/body/closed_call/checkpoint/concatenate" stack_frame_id=0} + %convert.324 = f32[1,128,8,128]{3,1,2,0:T(8,128)} convert(%pad.61.clone.3) diff --git a/tests/utils/sharding_info/deepseek2-16b/tpu7x-16/slice_1/rule_default/named_shardings.json b/tests/utils/sharding_info/deepseek2-16b/tpu7x-16/slice_1/rule_default/named_shardings.json index 4549ab46c5..b535731554 100644 --- a/tests/utils/sharding_info/deepseek2-16b/tpu7x-16/slice_1/rule_default/named_shardings.json +++ b/tests/utils/sharding_info/deepseek2-16b/tpu7x-16/slice_1/rule_default/named_shardings.json @@ -10,7 +10,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -24,7 +23,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -44,7 +42,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -58,17 +55,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ] + "tensor" ], "shape": [ 2048 @@ -85,7 +78,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -99,7 +91,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -108,7 +99,6 @@ "partition_spec": [ [ "fsdp", - "tensor_transpose", "context", "expert" ], @@ -137,7 +127,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -151,7 +140,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -160,7 +148,6 @@ "partition_spec": [ [ "fsdp", - "tensor_transpose", "context", "expert" ], @@ -189,7 +176,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -203,7 +189,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -219,7 +204,6 @@ null, [ "fsdp", - "tensor_transpose", "context", "expert" ] @@ -241,7 +225,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -255,17 +238,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ], + "tensor", null ], "shape": [ @@ -284,7 +263,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -298,17 +276,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ], + "tensor", null ], "shape": [ @@ -327,7 +301,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -341,17 +314,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ], + "tensor", null ], "shape": [ @@ -370,7 +339,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -384,7 +352,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -393,7 +360,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -424,7 +390,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -438,7 +403,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -454,7 +418,6 @@ null, [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -478,7 +441,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -492,7 +454,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -502,7 +463,6 @@ [ "fsdp", "fsdp_transpose", - "tensor_transpose", "context", "expert" ], @@ -526,7 +486,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -540,7 +499,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -556,7 +514,6 @@ null, [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -580,7 +537,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -594,7 +550,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -609,7 +564,6 @@ ], [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ] @@ -630,7 +584,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -644,7 +597,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -654,7 +606,6 @@ [ "fsdp", "fsdp_transpose", - "tensor_transpose", "context" ], null, @@ -677,7 +628,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -691,7 +641,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -702,7 +651,6 @@ null, [ "fsdp", - "tensor_transpose", "context" ], [ @@ -730,7 +678,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -744,7 +691,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -755,7 +701,6 @@ null, [ "fsdp", - "tensor_transpose", "context" ], [ @@ -783,7 +728,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -797,7 +741,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -814,7 +757,6 @@ ], [ "fsdp", - "tensor_transpose", "context" ] ], @@ -836,7 +778,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -850,7 +791,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -859,7 +799,6 @@ "partition_spec": [ [ "fsdp", - "tensor_transpose", "context", "expert" ], @@ -888,7 +827,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -902,7 +840,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -911,7 +848,6 @@ "partition_spec": [ [ "fsdp", - "tensor_transpose", "context", "expert" ], @@ -940,7 +876,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -954,7 +889,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -970,7 +904,6 @@ null, [ "fsdp", - "tensor_transpose", "context", "expert" ] @@ -992,7 +925,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1006,17 +938,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ], + "tensor", null ], "shape": [ @@ -1035,7 +963,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1049,17 +976,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ], + "tensor", null ], "shape": [ @@ -1078,7 +1001,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1092,17 +1014,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ], + "tensor", null ], "shape": [ @@ -1121,7 +1039,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1135,7 +1052,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -1144,7 +1060,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -1175,7 +1090,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1189,7 +1103,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -1205,7 +1118,6 @@ null, [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -1229,7 +1141,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1243,7 +1154,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -1253,7 +1163,6 @@ [ "fsdp", "fsdp_transpose", - "tensor_transpose", "context", "expert" ], @@ -1277,7 +1186,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1291,7 +1199,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -1307,7 +1214,6 @@ null, [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -1331,7 +1237,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1345,7 +1250,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -1354,7 +1258,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -1381,7 +1284,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1395,7 +1297,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -1415,7 +1316,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1429,17 +1329,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ] + "tensor" ], "shape": [ 2048 @@ -1456,7 +1352,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1470,7 +1365,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -1479,7 +1373,6 @@ "partition_spec": [ [ "fsdp", - "tensor_transpose", "context", "expert" ], @@ -1508,7 +1401,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1522,7 +1414,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -1531,7 +1422,6 @@ "partition_spec": [ [ "fsdp", - "tensor_transpose", "context", "expert" ], @@ -1560,7 +1450,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1574,7 +1463,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -1590,7 +1478,6 @@ null, [ "fsdp", - "tensor_transpose", "context", "expert" ] @@ -1612,7 +1499,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1626,17 +1512,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ], + "tensor", null ], "shape": [ @@ -1655,7 +1537,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1669,17 +1550,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ], + "tensor", null ], "shape": [ @@ -1698,7 +1575,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1712,17 +1588,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ], + "tensor", null ], "shape": [ @@ -1741,7 +1613,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1755,7 +1626,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -1764,7 +1634,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -1795,7 +1664,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1809,7 +1677,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -1825,7 +1692,6 @@ null, [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -1849,7 +1715,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1863,7 +1728,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -1873,7 +1737,6 @@ [ "fsdp", "fsdp_transpose", - "tensor_transpose", "context", "expert" ], @@ -1897,7 +1760,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1911,7 +1773,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -1927,7 +1788,6 @@ null, [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -1951,7 +1811,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1965,7 +1824,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -1980,7 +1838,6 @@ ], [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ] @@ -2001,7 +1858,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2015,7 +1871,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -2025,7 +1880,6 @@ [ "fsdp", "fsdp_transpose", - "tensor_transpose", "context" ], null, @@ -2048,7 +1902,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2062,7 +1915,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -2073,7 +1925,6 @@ null, [ "fsdp", - "tensor_transpose", "context" ], [ @@ -2101,7 +1952,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2115,7 +1965,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -2126,7 +1975,6 @@ null, [ "fsdp", - "tensor_transpose", "context" ], [ @@ -2154,7 +2002,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2168,7 +2015,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -2185,7 +2031,6 @@ ], [ "fsdp", - "tensor_transpose", "context" ] ], @@ -2207,7 +2052,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2221,7 +2065,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -2230,7 +2073,6 @@ "partition_spec": [ [ "fsdp", - "tensor_transpose", "context", "expert" ], @@ -2259,7 +2101,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2273,7 +2114,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -2282,7 +2122,6 @@ "partition_spec": [ [ "fsdp", - "tensor_transpose", "context", "expert" ], @@ -2311,7 +2150,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2325,7 +2163,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -2341,7 +2178,6 @@ null, [ "fsdp", - "tensor_transpose", "context", "expert" ] @@ -2363,7 +2199,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2377,17 +2212,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ], + "tensor", null ], "shape": [ @@ -2406,7 +2237,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2420,17 +2250,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ], + "tensor", null ], "shape": [ @@ -2449,7 +2275,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2463,17 +2288,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ], + "tensor", null ], "shape": [ @@ -2492,7 +2313,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2506,7 +2326,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -2515,7 +2334,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -2546,7 +2364,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2560,7 +2377,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -2576,7 +2392,6 @@ null, [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -2600,7 +2415,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2614,7 +2428,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -2624,7 +2437,6 @@ [ "fsdp", "fsdp_transpose", - "tensor_transpose", "context", "expert" ], @@ -2648,7 +2460,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2662,7 +2473,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -2678,7 +2488,6 @@ null, [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -2702,7 +2511,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2716,7 +2524,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -2725,7 +2532,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -2752,7 +2558,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2766,17 +2571,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ] + "tensor" ], "shape": [ 2048 @@ -2793,7 +2594,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2807,7 +2607,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -2816,7 +2615,6 @@ "partition_spec": [ [ "fsdp", - "tensor_transpose", "context", "expert" ], @@ -2845,7 +2643,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2859,7 +2656,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -2868,7 +2664,6 @@ "partition_spec": [ [ "fsdp", - "tensor_transpose", "context", "expert" ], @@ -2897,7 +2692,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2911,7 +2705,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -2927,7 +2720,6 @@ null, [ "fsdp", - "tensor_transpose", "context", "expert" ] @@ -2949,7 +2741,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2963,17 +2754,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ], + "tensor", null ], "shape": [ @@ -2992,7 +2779,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3006,17 +2792,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ], + "tensor", null ], "shape": [ @@ -3035,7 +2817,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3049,17 +2830,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ], + "tensor", null ], "shape": [ @@ -3078,7 +2855,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3092,7 +2868,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -3101,7 +2876,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -3132,7 +2906,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3146,7 +2919,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -3162,7 +2934,6 @@ null, [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -3186,7 +2957,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3200,7 +2970,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -3210,7 +2979,6 @@ [ "fsdp", "fsdp_transpose", - "tensor_transpose", "context", "expert" ], @@ -3234,7 +3002,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3248,7 +3015,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -3264,7 +3030,6 @@ null, [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -3288,7 +3053,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3302,7 +3066,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -3317,7 +3080,6 @@ ], [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ] @@ -3338,7 +3100,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3352,7 +3113,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -3362,7 +3122,6 @@ [ "fsdp", "fsdp_transpose", - "tensor_transpose", "context" ], null, @@ -3385,7 +3144,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3399,7 +3157,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -3410,7 +3167,6 @@ null, [ "fsdp", - "tensor_transpose", "context" ], [ @@ -3438,7 +3194,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3452,7 +3207,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -3463,7 +3217,6 @@ null, [ "fsdp", - "tensor_transpose", "context" ], [ @@ -3491,7 +3244,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3505,7 +3257,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -3522,7 +3273,6 @@ ], [ "fsdp", - "tensor_transpose", "context" ] ], @@ -3544,7 +3294,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3558,7 +3307,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -3567,7 +3315,6 @@ "partition_spec": [ [ "fsdp", - "tensor_transpose", "context", "expert" ], @@ -3596,7 +3343,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3610,7 +3356,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -3619,7 +3364,6 @@ "partition_spec": [ [ "fsdp", - "tensor_transpose", "context", "expert" ], @@ -3648,7 +3392,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3662,7 +3405,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -3678,7 +3420,6 @@ null, [ "fsdp", - "tensor_transpose", "context", "expert" ] @@ -3700,7 +3441,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3714,17 +3454,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ], + "tensor", null ], "shape": [ @@ -3743,7 +3479,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3757,17 +3492,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ], + "tensor", null ], "shape": [ @@ -3786,7 +3517,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3800,17 +3530,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ], + "tensor", null ], "shape": [ @@ -3829,7 +3555,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3843,7 +3568,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -3852,7 +3576,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -3883,7 +3606,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3897,7 +3619,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -3913,7 +3634,6 @@ null, [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -3937,7 +3657,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3951,7 +3670,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -3961,7 +3679,6 @@ [ "fsdp", "fsdp_transpose", - "tensor_transpose", "context", "expert" ], @@ -3985,7 +3702,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3999,7 +3715,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -4015,7 +3730,6 @@ null, [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -4039,7 +3753,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -4053,7 +3766,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -4062,7 +3774,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -4089,7 +3800,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -4103,7 +3813,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 diff --git a/tests/utils/sharding_info/deepseek2-16b/v6e-16/slice_1/rule_default_ici_fsdp_parallelism=-1_ici_expert_parallelism=4/named_shardings.json b/tests/utils/sharding_info/deepseek2-16b/v6e-16/slice_1/rule_default_ici_fsdp_parallelism=-1_ici_expert_parallelism=4/named_shardings.json index ca994b50f6..c6144b3769 100644 --- a/tests/utils/sharding_info/deepseek2-16b/v6e-16/slice_1/rule_default_ici_fsdp_parallelism=-1_ici_expert_parallelism=4/named_shardings.json +++ b/tests/utils/sharding_info/deepseek2-16b/v6e-16/slice_1/rule_default_ici_fsdp_parallelism=-1_ici_expert_parallelism=4/named_shardings.json @@ -10,7 +10,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -24,7 +23,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 @@ -44,7 +42,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -58,17 +55,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ] + "tensor" ], "shape": [ 2048 @@ -85,7 +78,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -99,7 +91,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 @@ -108,7 +99,6 @@ "partition_spec": [ [ "fsdp", - "tensor_transpose", "context", "expert" ], @@ -137,7 +127,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -151,7 +140,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 @@ -160,7 +148,6 @@ "partition_spec": [ [ "fsdp", - "tensor_transpose", "context", "expert" ], @@ -189,7 +176,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -203,7 +189,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 @@ -219,7 +204,6 @@ null, [ "fsdp", - "tensor_transpose", "context", "expert" ] @@ -241,7 +225,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -255,17 +238,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ], + "tensor", null ], "shape": [ @@ -284,7 +263,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -298,17 +276,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ], + "tensor", null ], "shape": [ @@ -327,7 +301,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -341,17 +314,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ], + "tensor", null ], "shape": [ @@ -370,7 +339,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -384,7 +352,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 @@ -393,7 +360,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -424,7 +390,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -438,7 +403,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 @@ -454,7 +418,6 @@ null, [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -478,7 +441,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -492,7 +454,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 @@ -502,7 +463,6 @@ [ "fsdp", "fsdp_transpose", - "tensor_transpose", "context", "expert" ], @@ -526,7 +486,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -540,7 +499,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 @@ -556,7 +514,6 @@ null, [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -580,7 +537,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -594,7 +550,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 @@ -609,7 +564,6 @@ ], [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ] @@ -630,7 +584,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -644,7 +597,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 @@ -654,7 +606,6 @@ [ "fsdp", "fsdp_transpose", - "tensor_transpose", "context" ], null, @@ -677,7 +628,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -691,7 +641,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 @@ -702,7 +651,6 @@ null, [ "fsdp", - "tensor_transpose", "context" ], [ @@ -730,7 +678,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -744,7 +691,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 @@ -755,7 +701,6 @@ null, [ "fsdp", - "tensor_transpose", "context" ], [ @@ -783,7 +728,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -797,7 +741,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 @@ -814,7 +757,6 @@ ], [ "fsdp", - "tensor_transpose", "context" ] ], @@ -836,7 +778,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -850,7 +791,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 @@ -859,7 +799,6 @@ "partition_spec": [ [ "fsdp", - "tensor_transpose", "context", "expert" ], @@ -888,7 +827,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -902,7 +840,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 @@ -911,7 +848,6 @@ "partition_spec": [ [ "fsdp", - "tensor_transpose", "context", "expert" ], @@ -940,7 +876,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -954,7 +889,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 @@ -970,7 +904,6 @@ null, [ "fsdp", - "tensor_transpose", "context", "expert" ] @@ -992,7 +925,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1006,17 +938,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ], + "tensor", null ], "shape": [ @@ -1035,7 +963,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1049,17 +976,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ], + "tensor", null ], "shape": [ @@ -1078,7 +1001,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1092,17 +1014,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ], + "tensor", null ], "shape": [ @@ -1121,7 +1039,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1135,7 +1052,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 @@ -1144,7 +1060,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -1175,7 +1090,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1189,7 +1103,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 @@ -1205,7 +1118,6 @@ null, [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -1229,7 +1141,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1243,7 +1154,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 @@ -1253,7 +1163,6 @@ [ "fsdp", "fsdp_transpose", - "tensor_transpose", "context", "expert" ], @@ -1277,7 +1186,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1291,7 +1199,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 @@ -1307,7 +1214,6 @@ null, [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -1331,7 +1237,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1345,7 +1250,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 @@ -1354,7 +1258,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -1381,7 +1284,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1395,7 +1297,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 @@ -1415,7 +1316,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1429,17 +1329,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ] + "tensor" ], "shape": [ 2048 @@ -1456,7 +1352,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1470,7 +1365,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 @@ -1479,7 +1373,6 @@ "partition_spec": [ [ "fsdp", - "tensor_transpose", "context", "expert" ], @@ -1508,7 +1401,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1522,7 +1414,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 @@ -1531,7 +1422,6 @@ "partition_spec": [ [ "fsdp", - "tensor_transpose", "context", "expert" ], @@ -1560,7 +1450,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1574,7 +1463,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 @@ -1590,7 +1478,6 @@ null, [ "fsdp", - "tensor_transpose", "context", "expert" ] @@ -1612,7 +1499,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1626,17 +1512,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ], + "tensor", null ], "shape": [ @@ -1655,7 +1537,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1669,17 +1550,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ], + "tensor", null ], "shape": [ @@ -1698,7 +1575,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1712,17 +1588,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ], + "tensor", null ], "shape": [ @@ -1741,7 +1613,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1755,7 +1626,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 @@ -1764,7 +1634,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -1795,7 +1664,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1809,7 +1677,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 @@ -1825,7 +1692,6 @@ null, [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -1849,7 +1715,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1863,7 +1728,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 @@ -1873,7 +1737,6 @@ [ "fsdp", "fsdp_transpose", - "tensor_transpose", "context", "expert" ], @@ -1897,7 +1760,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1911,7 +1773,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 @@ -1927,7 +1788,6 @@ null, [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -1951,7 +1811,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1965,7 +1824,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 @@ -1980,7 +1838,6 @@ ], [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ] @@ -2001,7 +1858,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2015,7 +1871,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 @@ -2025,7 +1880,6 @@ [ "fsdp", "fsdp_transpose", - "tensor_transpose", "context" ], null, @@ -2048,7 +1902,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2062,7 +1915,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 @@ -2073,7 +1925,6 @@ null, [ "fsdp", - "tensor_transpose", "context" ], [ @@ -2101,7 +1952,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2115,7 +1965,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 @@ -2126,7 +1975,6 @@ null, [ "fsdp", - "tensor_transpose", "context" ], [ @@ -2154,7 +2002,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2168,7 +2015,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 @@ -2185,7 +2031,6 @@ ], [ "fsdp", - "tensor_transpose", "context" ] ], @@ -2207,7 +2052,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2221,7 +2065,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 @@ -2230,7 +2073,6 @@ "partition_spec": [ [ "fsdp", - "tensor_transpose", "context", "expert" ], @@ -2259,7 +2101,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2273,7 +2114,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 @@ -2282,7 +2122,6 @@ "partition_spec": [ [ "fsdp", - "tensor_transpose", "context", "expert" ], @@ -2311,7 +2150,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2325,7 +2163,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 @@ -2341,7 +2178,6 @@ null, [ "fsdp", - "tensor_transpose", "context", "expert" ] @@ -2363,7 +2199,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2377,17 +2212,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ], + "tensor", null ], "shape": [ @@ -2406,7 +2237,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2420,17 +2250,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ], + "tensor", null ], "shape": [ @@ -2449,7 +2275,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2463,17 +2288,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ], + "tensor", null ], "shape": [ @@ -2492,7 +2313,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2506,7 +2326,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 @@ -2515,7 +2334,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -2546,7 +2364,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2560,7 +2377,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 @@ -2576,7 +2392,6 @@ null, [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -2600,7 +2415,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2614,7 +2428,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 @@ -2624,7 +2437,6 @@ [ "fsdp", "fsdp_transpose", - "tensor_transpose", "context", "expert" ], @@ -2648,7 +2460,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2662,7 +2473,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 @@ -2678,7 +2488,6 @@ null, [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -2702,7 +2511,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2716,7 +2524,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 @@ -2725,7 +2532,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -2752,7 +2558,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2766,17 +2571,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ] + "tensor" ], "shape": [ 2048 @@ -2793,7 +2594,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2807,7 +2607,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 @@ -2816,7 +2615,6 @@ "partition_spec": [ [ "fsdp", - "tensor_transpose", "context", "expert" ], @@ -2845,7 +2643,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2859,7 +2656,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 @@ -2868,7 +2664,6 @@ "partition_spec": [ [ "fsdp", - "tensor_transpose", "context", "expert" ], @@ -2897,7 +2692,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2911,7 +2705,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 @@ -2927,7 +2720,6 @@ null, [ "fsdp", - "tensor_transpose", "context", "expert" ] @@ -2949,7 +2741,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2963,17 +2754,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ], + "tensor", null ], "shape": [ @@ -2992,7 +2779,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3006,17 +2792,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ], + "tensor", null ], "shape": [ @@ -3035,7 +2817,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3049,17 +2830,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ], + "tensor", null ], "shape": [ @@ -3078,7 +2855,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3092,7 +2868,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 @@ -3101,7 +2876,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -3132,7 +2906,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3146,7 +2919,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 @@ -3162,7 +2934,6 @@ null, [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -3186,7 +2957,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3200,7 +2970,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 @@ -3210,7 +2979,6 @@ [ "fsdp", "fsdp_transpose", - "tensor_transpose", "context", "expert" ], @@ -3234,7 +3002,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3248,7 +3015,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 @@ -3264,7 +3030,6 @@ null, [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -3288,7 +3053,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3302,7 +3066,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 @@ -3317,7 +3080,6 @@ ], [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ] @@ -3338,7 +3100,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3352,7 +3113,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 @@ -3362,7 +3122,6 @@ [ "fsdp", "fsdp_transpose", - "tensor_transpose", "context" ], null, @@ -3385,7 +3144,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3399,7 +3157,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 @@ -3410,7 +3167,6 @@ null, [ "fsdp", - "tensor_transpose", "context" ], [ @@ -3438,7 +3194,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3452,7 +3207,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 @@ -3463,7 +3217,6 @@ null, [ "fsdp", - "tensor_transpose", "context" ], [ @@ -3491,7 +3244,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3505,7 +3257,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 @@ -3522,7 +3273,6 @@ ], [ "fsdp", - "tensor_transpose", "context" ] ], @@ -3544,7 +3294,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3558,7 +3307,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 @@ -3567,7 +3315,6 @@ "partition_spec": [ [ "fsdp", - "tensor_transpose", "context", "expert" ], @@ -3596,7 +3343,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3610,7 +3356,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 @@ -3619,7 +3364,6 @@ "partition_spec": [ [ "fsdp", - "tensor_transpose", "context", "expert" ], @@ -3648,7 +3392,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3662,7 +3405,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 @@ -3678,7 +3420,6 @@ null, [ "fsdp", - "tensor_transpose", "context", "expert" ] @@ -3700,7 +3441,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3714,17 +3454,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ], + "tensor", null ], "shape": [ @@ -3743,7 +3479,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3757,17 +3492,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ], + "tensor", null ], "shape": [ @@ -3786,7 +3517,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3800,17 +3530,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ], + "tensor", null ], "shape": [ @@ -3829,7 +3555,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3843,7 +3568,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 @@ -3852,7 +3576,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -3883,7 +3606,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3897,7 +3619,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 @@ -3913,7 +3634,6 @@ null, [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -3937,7 +3657,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3951,7 +3670,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 @@ -3961,7 +3679,6 @@ [ "fsdp", "fsdp_transpose", - "tensor_transpose", "context", "expert" ], @@ -3985,7 +3702,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3999,7 +3715,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 @@ -4015,7 +3730,6 @@ null, [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -4039,7 +3753,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -4053,7 +3766,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 @@ -4062,7 +3774,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -4089,7 +3800,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -4103,7 +3813,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 4, "autoregressive": 1 diff --git a/tests/utils/sharding_info/gpt-oss-20b/tpu7x-16/slice_1/rule_default/named_shardings.json b/tests/utils/sharding_info/gpt-oss-20b/tpu7x-16/slice_1/rule_default/named_shardings.json index 7e1b2785ae..970aa1d4ef 100644 --- a/tests/utils/sharding_info/gpt-oss-20b/tpu7x-16/slice_1/rule_default/named_shardings.json +++ b/tests/utils/sharding_info/gpt-oss-20b/tpu7x-16/slice_1/rule_default/named_shardings.json @@ -10,7 +10,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -24,7 +23,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -44,7 +42,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -58,17 +55,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ] + "tensor" ], "shape": [ 2880 @@ -85,7 +78,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -99,7 +91,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -108,7 +99,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -132,7 +122,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -146,7 +135,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -162,7 +150,6 @@ "stage", [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -186,7 +173,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -200,7 +186,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -210,7 +195,6 @@ [ "fsdp", "fsdp_transpose", - "tensor_transpose", "context", "expert" ], @@ -232,7 +216,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -246,7 +229,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -255,7 +237,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -286,7 +267,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -300,7 +280,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -309,7 +288,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -333,7 +311,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -347,7 +324,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -363,7 +339,6 @@ "stage", [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -387,7 +362,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -401,7 +375,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -427,7 +400,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -441,7 +413,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -450,7 +421,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -474,7 +444,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -488,7 +457,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -504,7 +472,6 @@ "stage", [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -528,7 +495,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -542,7 +508,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -568,7 +533,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -582,7 +546,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -592,7 +555,6 @@ [ "fsdp", "fsdp_transpose", - "tensor_transpose", "context", "expert" ], @@ -616,7 +578,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -630,7 +591,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -641,7 +601,6 @@ "stage", [ "fsdp", - "tensor_transpose", "context" ], [ @@ -669,7 +628,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -683,7 +641,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -694,7 +651,6 @@ "stage", [ "tensor", - "tensor_transpose", "tensor_sequence" ] ], @@ -715,7 +671,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -729,7 +684,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -740,7 +694,6 @@ "stage", [ "fsdp", - "tensor_transpose", "context" ], [ @@ -768,7 +721,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -782,7 +734,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -793,7 +744,6 @@ "stage", [ "tensor", - "tensor_transpose", "tensor_sequence" ] ], @@ -814,7 +764,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -828,7 +777,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -845,7 +793,6 @@ ], [ "fsdp", - "tensor_transpose", "context" ] ], @@ -867,7 +814,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -881,7 +827,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -890,10 +835,7 @@ "partition_spec": [ "expert", "stage", - [ - "tensor", - "tensor_transpose" - ] + "tensor" ], "shape": [ 32, @@ -912,7 +854,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -926,17 +867,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ], + "tensor", "stage" ], "shape": [ @@ -955,7 +892,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -969,17 +905,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ], + "tensor", "stage" ], "shape": [ @@ -998,7 +930,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1012,7 +943,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -1021,7 +951,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -1045,7 +974,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1059,7 +987,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -1075,7 +1002,6 @@ "stage", [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -1099,7 +1025,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1113,7 +1038,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -1123,7 +1047,6 @@ [ "fsdp", "fsdp_transpose", - "tensor_transpose", "context", "expert" ], @@ -1145,7 +1068,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1159,7 +1081,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -1168,7 +1089,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -1199,7 +1119,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1213,7 +1132,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -1222,7 +1140,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -1246,7 +1163,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1260,7 +1176,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -1276,7 +1191,6 @@ "stage", [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -1300,7 +1214,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1314,7 +1227,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -1340,7 +1252,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1354,7 +1265,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -1363,7 +1273,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -1387,7 +1296,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1401,7 +1309,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -1417,7 +1324,6 @@ "stage", [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -1441,7 +1347,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1455,7 +1360,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -1481,7 +1385,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1495,7 +1398,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -1505,7 +1407,6 @@ [ "fsdp", "fsdp_transpose", - "tensor_transpose", "context", "expert" ], @@ -1529,7 +1430,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1543,7 +1443,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -1554,7 +1453,6 @@ "stage", [ "fsdp", - "tensor_transpose", "context" ], [ @@ -1582,7 +1480,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1596,7 +1493,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -1607,7 +1503,6 @@ "stage", [ "tensor", - "tensor_transpose", "tensor_sequence" ] ], @@ -1628,7 +1523,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1642,7 +1536,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -1653,7 +1546,6 @@ "stage", [ "fsdp", - "tensor_transpose", "context" ], [ @@ -1681,7 +1573,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1695,7 +1586,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -1706,7 +1596,6 @@ "stage", [ "tensor", - "tensor_transpose", "tensor_sequence" ] ], @@ -1727,7 +1616,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1741,7 +1629,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -1758,7 +1645,6 @@ ], [ "fsdp", - "tensor_transpose", "context" ] ], @@ -1780,7 +1666,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1794,7 +1679,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -1803,10 +1687,7 @@ "partition_spec": [ "expert", "stage", - [ - "tensor", - "tensor_transpose" - ] + "tensor" ], "shape": [ 32, @@ -1825,7 +1706,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1839,17 +1719,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ], + "tensor", "stage" ], "shape": [ @@ -1868,7 +1744,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1882,17 +1757,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ], + "tensor", "stage" ], "shape": [ @@ -1911,7 +1782,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1925,7 +1795,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -1940,7 +1809,6 @@ ], [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ] @@ -1961,7 +1829,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1975,7 +1842,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -1984,7 +1850,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -2011,7 +1876,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2025,7 +1889,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -2045,7 +1908,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2059,17 +1921,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ] + "tensor" ], "shape": [ 2880 @@ -2086,7 +1944,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2100,7 +1957,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -2109,7 +1965,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -2133,7 +1988,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2147,7 +2001,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -2163,7 +2016,6 @@ "stage", [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -2187,7 +2039,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2201,7 +2052,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -2211,7 +2061,6 @@ [ "fsdp", "fsdp_transpose", - "tensor_transpose", "context", "expert" ], @@ -2233,7 +2082,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2247,7 +2095,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -2256,7 +2103,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -2287,7 +2133,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2301,7 +2146,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -2310,7 +2154,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -2334,7 +2177,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2348,7 +2190,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -2364,7 +2205,6 @@ "stage", [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -2388,7 +2228,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2402,7 +2241,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -2428,7 +2266,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2442,7 +2279,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -2451,7 +2287,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -2475,7 +2310,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2489,7 +2323,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -2505,7 +2338,6 @@ "stage", [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -2529,7 +2361,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2543,7 +2374,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -2569,7 +2399,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2583,7 +2412,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -2593,7 +2421,6 @@ [ "fsdp", "fsdp_transpose", - "tensor_transpose", "context", "expert" ], @@ -2617,7 +2444,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2631,7 +2457,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -2642,7 +2467,6 @@ "stage", [ "fsdp", - "tensor_transpose", "context" ], [ @@ -2670,7 +2494,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2684,7 +2507,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -2695,7 +2517,6 @@ "stage", [ "tensor", - "tensor_transpose", "tensor_sequence" ] ], @@ -2716,7 +2537,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2730,7 +2550,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -2741,7 +2560,6 @@ "stage", [ "fsdp", - "tensor_transpose", "context" ], [ @@ -2769,7 +2587,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2783,7 +2600,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -2794,7 +2610,6 @@ "stage", [ "tensor", - "tensor_transpose", "tensor_sequence" ] ], @@ -2815,7 +2630,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2829,7 +2643,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -2846,7 +2659,6 @@ ], [ "fsdp", - "tensor_transpose", "context" ] ], @@ -2868,7 +2680,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2882,7 +2693,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -2891,10 +2701,7 @@ "partition_spec": [ "expert", "stage", - [ - "tensor", - "tensor_transpose" - ] + "tensor" ], "shape": [ 32, @@ -2913,7 +2720,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2927,17 +2733,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ], + "tensor", "stage" ], "shape": [ @@ -2956,7 +2758,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2970,17 +2771,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ], + "tensor", "stage" ], "shape": [ @@ -2999,7 +2796,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3013,7 +2809,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -3022,7 +2817,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -3046,7 +2840,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3060,7 +2853,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -3076,7 +2868,6 @@ "stage", [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -3100,7 +2891,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3114,7 +2904,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -3124,7 +2913,6 @@ [ "fsdp", "fsdp_transpose", - "tensor_transpose", "context", "expert" ], @@ -3146,7 +2934,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3160,7 +2947,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -3169,7 +2955,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -3200,7 +2985,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3214,7 +2998,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -3223,7 +3006,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -3247,7 +3029,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3261,7 +3042,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -3277,7 +3057,6 @@ "stage", [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -3301,7 +3080,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3315,7 +3093,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -3341,7 +3118,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3355,7 +3131,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -3364,7 +3139,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -3388,7 +3162,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3402,7 +3175,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -3418,7 +3190,6 @@ "stage", [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -3442,7 +3213,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3456,7 +3226,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -3482,7 +3251,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3496,7 +3264,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -3506,7 +3273,6 @@ [ "fsdp", "fsdp_transpose", - "tensor_transpose", "context", "expert" ], @@ -3530,7 +3296,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3544,7 +3309,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -3555,7 +3319,6 @@ "stage", [ "fsdp", - "tensor_transpose", "context" ], [ @@ -3583,7 +3346,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3597,7 +3359,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -3608,7 +3369,6 @@ "stage", [ "tensor", - "tensor_transpose", "tensor_sequence" ] ], @@ -3629,7 +3389,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3643,7 +3402,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -3654,7 +3412,6 @@ "stage", [ "fsdp", - "tensor_transpose", "context" ], [ @@ -3682,7 +3439,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3696,7 +3452,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -3707,7 +3462,6 @@ "stage", [ "tensor", - "tensor_transpose", "tensor_sequence" ] ], @@ -3728,7 +3482,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3742,7 +3495,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -3759,7 +3511,6 @@ ], [ "fsdp", - "tensor_transpose", "context" ] ], @@ -3781,7 +3532,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3795,7 +3545,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -3804,10 +3553,7 @@ "partition_spec": [ "expert", "stage", - [ - "tensor", - "tensor_transpose" - ] + "tensor" ], "shape": [ 32, @@ -3826,7 +3572,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3840,17 +3585,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ], + "tensor", "stage" ], "shape": [ @@ -3869,7 +3610,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3883,17 +3623,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ], + "tensor", "stage" ], "shape": [ @@ -3912,7 +3648,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3926,7 +3661,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -3941,7 +3675,6 @@ ], [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ] @@ -3962,7 +3695,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3976,7 +3708,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -3985,7 +3716,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -4012,7 +3742,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -4026,17 +3755,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ] + "tensor" ], "shape": [ 2880 @@ -4053,7 +3778,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -4067,7 +3791,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -4076,7 +3799,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -4100,7 +3822,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -4114,7 +3835,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -4130,7 +3850,6 @@ "stage", [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -4154,7 +3873,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -4168,7 +3886,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -4178,7 +3895,6 @@ [ "fsdp", "fsdp_transpose", - "tensor_transpose", "context", "expert" ], @@ -4200,7 +3916,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -4214,7 +3929,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -4223,7 +3937,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -4254,7 +3967,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -4268,7 +3980,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -4277,7 +3988,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -4301,7 +4011,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -4315,7 +4024,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -4331,7 +4039,6 @@ "stage", [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -4355,7 +4062,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -4369,7 +4075,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -4395,7 +4100,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -4409,7 +4113,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -4418,7 +4121,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -4442,7 +4144,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -4456,7 +4157,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -4472,7 +4172,6 @@ "stage", [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -4496,7 +4195,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -4510,7 +4208,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -4536,7 +4233,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -4550,7 +4246,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -4560,7 +4255,6 @@ [ "fsdp", "fsdp_transpose", - "tensor_transpose", "context", "expert" ], @@ -4584,7 +4278,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -4598,7 +4291,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -4609,7 +4301,6 @@ "stage", [ "fsdp", - "tensor_transpose", "context" ], [ @@ -4637,7 +4328,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -4651,7 +4341,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -4662,7 +4351,6 @@ "stage", [ "tensor", - "tensor_transpose", "tensor_sequence" ] ], @@ -4683,7 +4371,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -4697,7 +4384,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -4708,7 +4394,6 @@ "stage", [ "fsdp", - "tensor_transpose", "context" ], [ @@ -4736,7 +4421,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -4750,7 +4434,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -4761,7 +4444,6 @@ "stage", [ "tensor", - "tensor_transpose", "tensor_sequence" ] ], @@ -4782,7 +4464,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -4796,7 +4477,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -4813,7 +4493,6 @@ ], [ "fsdp", - "tensor_transpose", "context" ] ], @@ -4835,7 +4514,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -4849,7 +4527,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -4858,10 +4535,7 @@ "partition_spec": [ "expert", "stage", - [ - "tensor", - "tensor_transpose" - ] + "tensor" ], "shape": [ 32, @@ -4880,7 +4554,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -4894,17 +4567,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ], + "tensor", "stage" ], "shape": [ @@ -4923,7 +4592,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -4937,17 +4605,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ], + "tensor", "stage" ], "shape": [ @@ -4966,7 +4630,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -4980,7 +4643,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -4989,7 +4651,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -5013,7 +4674,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -5027,7 +4687,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -5043,7 +4702,6 @@ "stage", [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -5067,7 +4725,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -5081,7 +4738,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -5091,7 +4747,6 @@ [ "fsdp", "fsdp_transpose", - "tensor_transpose", "context", "expert" ], @@ -5113,7 +4768,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -5127,7 +4781,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -5136,7 +4789,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -5167,7 +4819,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -5181,7 +4832,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -5190,7 +4840,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -5214,7 +4863,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -5228,7 +4876,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -5244,7 +4891,6 @@ "stage", [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -5268,7 +4914,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -5282,7 +4927,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -5308,7 +4952,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -5322,7 +4965,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -5331,7 +4973,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -5355,7 +4996,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -5369,7 +5009,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -5385,7 +5024,6 @@ "stage", [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -5409,7 +5047,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -5423,7 +5060,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -5449,7 +5085,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -5463,7 +5098,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -5473,7 +5107,6 @@ [ "fsdp", "fsdp_transpose", - "tensor_transpose", "context", "expert" ], @@ -5497,7 +5130,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -5511,7 +5143,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -5522,7 +5153,6 @@ "stage", [ "fsdp", - "tensor_transpose", "context" ], [ @@ -5550,7 +5180,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -5564,7 +5193,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -5575,7 +5203,6 @@ "stage", [ "tensor", - "tensor_transpose", "tensor_sequence" ] ], @@ -5596,7 +5223,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -5610,7 +5236,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -5621,7 +5246,6 @@ "stage", [ "fsdp", - "tensor_transpose", "context" ], [ @@ -5649,7 +5273,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -5663,7 +5286,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -5674,7 +5296,6 @@ "stage", [ "tensor", - "tensor_transpose", "tensor_sequence" ] ], @@ -5695,7 +5316,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -5709,7 +5329,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -5726,7 +5345,6 @@ ], [ "fsdp", - "tensor_transpose", "context" ] ], @@ -5748,7 +5366,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -5762,7 +5379,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -5771,10 +5387,7 @@ "partition_spec": [ "expert", "stage", - [ - "tensor", - "tensor_transpose" - ] + "tensor" ], "shape": [ 32, @@ -5793,7 +5406,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -5807,17 +5419,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ], + "tensor", "stage" ], "shape": [ @@ -5836,7 +5444,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -5850,17 +5457,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ], + "tensor", "stage" ], "shape": [ @@ -5879,7 +5482,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -5893,7 +5495,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -5908,7 +5509,6 @@ ], [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ] @@ -5929,7 +5529,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -5943,7 +5542,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -5952,7 +5550,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -5979,7 +5576,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -5993,7 +5589,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 diff --git a/tests/utils/sharding_info/gpt-oss-20b/tpu7x-16/slice_1/rule_default_ici_fsdp_parallelism=-1_ici_expert_parallelism=2/named_shardings.json b/tests/utils/sharding_info/gpt-oss-20b/tpu7x-16/slice_1/rule_default_ici_fsdp_parallelism=-1_ici_expert_parallelism=2/named_shardings.json index b3a0c7d967..1a90c3e31f 100644 --- a/tests/utils/sharding_info/gpt-oss-20b/tpu7x-16/slice_1/rule_default_ici_fsdp_parallelism=-1_ici_expert_parallelism=2/named_shardings.json +++ b/tests/utils/sharding_info/gpt-oss-20b/tpu7x-16/slice_1/rule_default_ici_fsdp_parallelism=-1_ici_expert_parallelism=2/named_shardings.json @@ -10,7 +10,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -24,7 +23,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -44,7 +42,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -58,17 +55,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ] + "tensor" ], "shape": [ 2880 @@ -85,7 +78,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -99,7 +91,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -108,7 +99,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -132,7 +122,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -146,7 +135,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -162,7 +150,6 @@ "stage", [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -186,7 +173,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -200,7 +186,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -210,7 +195,6 @@ [ "fsdp", "fsdp_transpose", - "tensor_transpose", "context", "expert" ], @@ -232,7 +216,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -246,7 +229,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -255,7 +237,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -286,7 +267,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -300,7 +280,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -309,7 +288,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -333,7 +311,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -347,7 +324,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -363,7 +339,6 @@ "stage", [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -387,7 +362,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -401,7 +375,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -427,7 +400,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -441,7 +413,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -450,7 +421,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -474,7 +444,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -488,7 +457,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -504,7 +472,6 @@ "stage", [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -528,7 +495,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -542,7 +508,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -568,7 +533,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -582,7 +546,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -592,7 +555,6 @@ [ "fsdp", "fsdp_transpose", - "tensor_transpose", "context", "expert" ], @@ -616,7 +578,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -630,7 +591,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -641,7 +601,6 @@ "stage", [ "fsdp", - "tensor_transpose", "context" ], [ @@ -669,7 +628,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -683,7 +641,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -694,7 +651,6 @@ "stage", [ "tensor", - "tensor_transpose", "tensor_sequence" ] ], @@ -715,7 +671,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -729,7 +684,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -740,7 +694,6 @@ "stage", [ "fsdp", - "tensor_transpose", "context" ], [ @@ -768,7 +721,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -782,7 +734,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -793,7 +744,6 @@ "stage", [ "tensor", - "tensor_transpose", "tensor_sequence" ] ], @@ -814,7 +764,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -828,7 +777,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -845,7 +793,6 @@ ], [ "fsdp", - "tensor_transpose", "context" ] ], @@ -867,7 +814,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -881,7 +827,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -890,10 +835,7 @@ "partition_spec": [ "expert", "stage", - [ - "tensor", - "tensor_transpose" - ] + "tensor" ], "shape": [ 32, @@ -912,7 +854,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -926,17 +867,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ], + "tensor", "stage" ], "shape": [ @@ -955,7 +892,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -969,17 +905,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ], + "tensor", "stage" ], "shape": [ @@ -998,7 +930,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1012,7 +943,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -1021,7 +951,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -1045,7 +974,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1059,7 +987,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -1075,7 +1002,6 @@ "stage", [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -1099,7 +1025,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1113,7 +1038,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -1123,7 +1047,6 @@ [ "fsdp", "fsdp_transpose", - "tensor_transpose", "context", "expert" ], @@ -1145,7 +1068,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1159,7 +1081,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -1168,7 +1089,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -1199,7 +1119,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1213,7 +1132,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -1222,7 +1140,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -1246,7 +1163,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1260,7 +1176,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -1276,7 +1191,6 @@ "stage", [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -1300,7 +1214,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1314,7 +1227,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -1340,7 +1252,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1354,7 +1265,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -1363,7 +1273,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -1387,7 +1296,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1401,7 +1309,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -1417,7 +1324,6 @@ "stage", [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -1441,7 +1347,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1455,7 +1360,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -1481,7 +1385,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1495,7 +1398,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -1505,7 +1407,6 @@ [ "fsdp", "fsdp_transpose", - "tensor_transpose", "context", "expert" ], @@ -1529,7 +1430,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1543,7 +1443,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -1554,7 +1453,6 @@ "stage", [ "fsdp", - "tensor_transpose", "context" ], [ @@ -1582,7 +1480,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1596,7 +1493,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -1607,7 +1503,6 @@ "stage", [ "tensor", - "tensor_transpose", "tensor_sequence" ] ], @@ -1628,7 +1523,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1642,7 +1536,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -1653,7 +1546,6 @@ "stage", [ "fsdp", - "tensor_transpose", "context" ], [ @@ -1681,7 +1573,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1695,7 +1586,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -1706,7 +1596,6 @@ "stage", [ "tensor", - "tensor_transpose", "tensor_sequence" ] ], @@ -1727,7 +1616,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1741,7 +1629,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -1758,7 +1645,6 @@ ], [ "fsdp", - "tensor_transpose", "context" ] ], @@ -1780,7 +1666,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1794,7 +1679,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -1803,10 +1687,7 @@ "partition_spec": [ "expert", "stage", - [ - "tensor", - "tensor_transpose" - ] + "tensor" ], "shape": [ 32, @@ -1825,7 +1706,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1839,17 +1719,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ], + "tensor", "stage" ], "shape": [ @@ -1868,7 +1744,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1882,17 +1757,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ], + "tensor", "stage" ], "shape": [ @@ -1911,7 +1782,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1925,7 +1795,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -1940,7 +1809,6 @@ ], [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ] @@ -1961,7 +1829,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1975,7 +1842,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -1984,7 +1850,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -2011,7 +1876,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2025,7 +1889,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -2045,7 +1908,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2059,17 +1921,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ] + "tensor" ], "shape": [ 2880 @@ -2086,7 +1944,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2100,7 +1957,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -2109,7 +1965,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -2133,7 +1988,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2147,7 +2001,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -2163,7 +2016,6 @@ "stage", [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -2187,7 +2039,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2201,7 +2052,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -2211,7 +2061,6 @@ [ "fsdp", "fsdp_transpose", - "tensor_transpose", "context", "expert" ], @@ -2233,7 +2082,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2247,7 +2095,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -2256,7 +2103,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -2287,7 +2133,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2301,7 +2146,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -2310,7 +2154,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -2334,7 +2177,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2348,7 +2190,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -2364,7 +2205,6 @@ "stage", [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -2388,7 +2228,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2402,7 +2241,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -2428,7 +2266,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2442,7 +2279,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -2451,7 +2287,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -2475,7 +2310,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2489,7 +2323,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -2505,7 +2338,6 @@ "stage", [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -2529,7 +2361,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2543,7 +2374,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -2569,7 +2399,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2583,7 +2412,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -2593,7 +2421,6 @@ [ "fsdp", "fsdp_transpose", - "tensor_transpose", "context", "expert" ], @@ -2617,7 +2444,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2631,7 +2457,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -2642,7 +2467,6 @@ "stage", [ "fsdp", - "tensor_transpose", "context" ], [ @@ -2670,7 +2494,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2684,7 +2507,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -2695,7 +2517,6 @@ "stage", [ "tensor", - "tensor_transpose", "tensor_sequence" ] ], @@ -2716,7 +2537,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2730,7 +2550,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -2741,7 +2560,6 @@ "stage", [ "fsdp", - "tensor_transpose", "context" ], [ @@ -2769,7 +2587,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2783,7 +2600,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -2794,7 +2610,6 @@ "stage", [ "tensor", - "tensor_transpose", "tensor_sequence" ] ], @@ -2815,7 +2630,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2829,7 +2643,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -2846,7 +2659,6 @@ ], [ "fsdp", - "tensor_transpose", "context" ] ], @@ -2868,7 +2680,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2882,7 +2693,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -2891,10 +2701,7 @@ "partition_spec": [ "expert", "stage", - [ - "tensor", - "tensor_transpose" - ] + "tensor" ], "shape": [ 32, @@ -2913,7 +2720,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2927,17 +2733,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ], + "tensor", "stage" ], "shape": [ @@ -2956,7 +2758,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -2970,17 +2771,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ], + "tensor", "stage" ], "shape": [ @@ -2999,7 +2796,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3013,7 +2809,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -3022,7 +2817,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -3046,7 +2840,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3060,7 +2853,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -3076,7 +2868,6 @@ "stage", [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -3100,7 +2891,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3114,7 +2904,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -3124,7 +2913,6 @@ [ "fsdp", "fsdp_transpose", - "tensor_transpose", "context", "expert" ], @@ -3146,7 +2934,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3160,7 +2947,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -3169,7 +2955,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -3200,7 +2985,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3214,7 +2998,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -3223,7 +3006,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -3247,7 +3029,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3261,7 +3042,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -3277,7 +3057,6 @@ "stage", [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -3301,7 +3080,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3315,7 +3093,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -3341,7 +3118,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3355,7 +3131,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -3364,7 +3139,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -3388,7 +3162,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3402,7 +3175,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -3418,7 +3190,6 @@ "stage", [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -3442,7 +3213,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3456,7 +3226,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -3482,7 +3251,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3496,7 +3264,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -3506,7 +3273,6 @@ [ "fsdp", "fsdp_transpose", - "tensor_transpose", "context", "expert" ], @@ -3530,7 +3296,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3544,7 +3309,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -3555,7 +3319,6 @@ "stage", [ "fsdp", - "tensor_transpose", "context" ], [ @@ -3583,7 +3346,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3597,7 +3359,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -3608,7 +3369,6 @@ "stage", [ "tensor", - "tensor_transpose", "tensor_sequence" ] ], @@ -3629,7 +3389,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3643,7 +3402,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -3654,7 +3412,6 @@ "stage", [ "fsdp", - "tensor_transpose", "context" ], [ @@ -3682,7 +3439,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3696,7 +3452,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -3707,7 +3462,6 @@ "stage", [ "tensor", - "tensor_transpose", "tensor_sequence" ] ], @@ -3728,7 +3482,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3742,7 +3495,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -3759,7 +3511,6 @@ ], [ "fsdp", - "tensor_transpose", "context" ] ], @@ -3781,7 +3532,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3795,7 +3545,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -3804,10 +3553,7 @@ "partition_spec": [ "expert", "stage", - [ - "tensor", - "tensor_transpose" - ] + "tensor" ], "shape": [ 32, @@ -3826,7 +3572,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3840,17 +3585,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ], + "tensor", "stage" ], "shape": [ @@ -3869,7 +3610,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3883,17 +3623,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ], + "tensor", "stage" ], "shape": [ @@ -3912,7 +3648,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3926,7 +3661,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -3941,7 +3675,6 @@ ], [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ] @@ -3962,7 +3695,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -3976,7 +3708,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -3985,7 +3716,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -4012,7 +3742,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -4026,17 +3755,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ] + "tensor" ], "shape": [ 2880 @@ -4053,7 +3778,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -4067,7 +3791,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -4076,7 +3799,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -4100,7 +3822,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -4114,7 +3835,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -4130,7 +3850,6 @@ "stage", [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -4154,7 +3873,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -4168,7 +3886,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -4178,7 +3895,6 @@ [ "fsdp", "fsdp_transpose", - "tensor_transpose", "context", "expert" ], @@ -4200,7 +3916,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -4214,7 +3929,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -4223,7 +3937,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -4254,7 +3967,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -4268,7 +3980,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -4277,7 +3988,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -4301,7 +4011,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -4315,7 +4024,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -4331,7 +4039,6 @@ "stage", [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -4355,7 +4062,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -4369,7 +4075,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -4395,7 +4100,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -4409,7 +4113,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -4418,7 +4121,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -4442,7 +4144,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -4456,7 +4157,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -4472,7 +4172,6 @@ "stage", [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -4496,7 +4195,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -4510,7 +4208,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -4536,7 +4233,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -4550,7 +4246,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -4560,7 +4255,6 @@ [ "fsdp", "fsdp_transpose", - "tensor_transpose", "context", "expert" ], @@ -4584,7 +4278,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -4598,7 +4291,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -4609,7 +4301,6 @@ "stage", [ "fsdp", - "tensor_transpose", "context" ], [ @@ -4637,7 +4328,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -4651,7 +4341,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -4662,7 +4351,6 @@ "stage", [ "tensor", - "tensor_transpose", "tensor_sequence" ] ], @@ -4683,7 +4371,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -4697,7 +4384,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -4708,7 +4394,6 @@ "stage", [ "fsdp", - "tensor_transpose", "context" ], [ @@ -4736,7 +4421,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -4750,7 +4434,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -4761,7 +4444,6 @@ "stage", [ "tensor", - "tensor_transpose", "tensor_sequence" ] ], @@ -4782,7 +4464,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -4796,7 +4477,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -4813,7 +4493,6 @@ ], [ "fsdp", - "tensor_transpose", "context" ] ], @@ -4835,7 +4514,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -4849,7 +4527,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -4858,10 +4535,7 @@ "partition_spec": [ "expert", "stage", - [ - "tensor", - "tensor_transpose" - ] + "tensor" ], "shape": [ 32, @@ -4880,7 +4554,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -4894,17 +4567,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ], + "tensor", "stage" ], "shape": [ @@ -4923,7 +4592,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -4937,17 +4605,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ], + "tensor", "stage" ], "shape": [ @@ -4966,7 +4630,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -4980,7 +4643,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -4989,7 +4651,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -5013,7 +4674,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -5027,7 +4687,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -5043,7 +4702,6 @@ "stage", [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -5067,7 +4725,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -5081,7 +4738,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -5091,7 +4747,6 @@ [ "fsdp", "fsdp_transpose", - "tensor_transpose", "context", "expert" ], @@ -5113,7 +4768,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -5127,7 +4781,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -5136,7 +4789,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -5167,7 +4819,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -5181,7 +4832,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -5190,7 +4840,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -5214,7 +4863,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -5228,7 +4876,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -5244,7 +4891,6 @@ "stage", [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -5268,7 +4914,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -5282,7 +4927,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -5308,7 +4952,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -5322,7 +4965,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -5331,7 +4973,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -5355,7 +4996,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -5369,7 +5009,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -5385,7 +5024,6 @@ "stage", [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -5409,7 +5047,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -5423,7 +5060,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -5449,7 +5085,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -5463,7 +5098,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -5473,7 +5107,6 @@ [ "fsdp", "fsdp_transpose", - "tensor_transpose", "context", "expert" ], @@ -5497,7 +5130,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -5511,7 +5143,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -5522,7 +5153,6 @@ "stage", [ "fsdp", - "tensor_transpose", "context" ], [ @@ -5550,7 +5180,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -5564,7 +5193,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -5575,7 +5203,6 @@ "stage", [ "tensor", - "tensor_transpose", "tensor_sequence" ] ], @@ -5596,7 +5223,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -5610,7 +5236,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -5621,7 +5246,6 @@ "stage", [ "fsdp", - "tensor_transpose", "context" ], [ @@ -5649,7 +5273,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -5663,7 +5286,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -5674,7 +5296,6 @@ "stage", [ "tensor", - "tensor_transpose", "tensor_sequence" ] ], @@ -5695,7 +5316,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -5709,7 +5329,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -5726,7 +5345,6 @@ ], [ "fsdp", - "tensor_transpose", "context" ] ], @@ -5748,7 +5366,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -5762,7 +5379,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -5771,10 +5387,7 @@ "partition_spec": [ "expert", "stage", - [ - "tensor", - "tensor_transpose" - ] + "tensor" ], "shape": [ 32, @@ -5793,7 +5406,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -5807,17 +5419,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ], + "tensor", "stage" ], "shape": [ @@ -5836,7 +5444,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -5850,17 +5457,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ], + "tensor", "stage" ], "shape": [ @@ -5879,7 +5482,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -5893,7 +5495,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -5908,7 +5509,6 @@ ], [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ] @@ -5929,7 +5529,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -5943,7 +5542,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 @@ -5952,7 +5550,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -5979,7 +5576,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -5993,7 +5589,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 2, "autoregressive": 1 diff --git a/tests/utils/sharding_info/qwen3-0.6b/tpu7x-16/slice_1/rule_default/named_shardings.json b/tests/utils/sharding_info/qwen3-0.6b/tpu7x-16/slice_1/rule_default/named_shardings.json index b86911af98..3c56796fba 100644 --- a/tests/utils/sharding_info/qwen3-0.6b/tpu7x-16/slice_1/rule_default/named_shardings.json +++ b/tests/utils/sharding_info/qwen3-0.6b/tpu7x-16/slice_1/rule_default/named_shardings.json @@ -10,7 +10,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -24,7 +23,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -44,7 +42,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -58,17 +55,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ] + "tensor" ], "shape": [ 1024 @@ -85,7 +78,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -99,7 +91,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -108,7 +99,6 @@ "partition_spec": [ [ "fsdp", - "tensor_transpose", "context", "expert" ], @@ -137,7 +127,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -151,7 +140,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -160,7 +148,6 @@ "partition_spec": [ [ "fsdp", - "tensor_transpose", "context", "expert" ], @@ -189,7 +176,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -203,7 +189,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -219,7 +204,6 @@ "stage", [ "fsdp", - "tensor_transpose", "context", "expert" ] @@ -241,7 +225,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -255,17 +238,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ], + "tensor", "stage" ], "shape": [ @@ -284,7 +263,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -298,17 +276,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ], + "tensor", "stage" ], "shape": [ @@ -327,7 +301,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -341,7 +314,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -357,7 +329,6 @@ "stage", [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -381,7 +352,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -395,17 +365,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ], + "tensor", "stage" ], "shape": [ @@ -424,7 +390,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -438,7 +403,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -447,7 +411,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -478,7 +441,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -492,7 +454,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -508,7 +469,6 @@ "stage", [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -532,7 +492,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -546,17 +505,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ], + "tensor", "stage" ], "shape": [ @@ -575,7 +530,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -589,7 +543,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -605,7 +558,6 @@ "stage", [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -629,7 +581,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -643,7 +594,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -652,7 +602,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -679,7 +628,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -693,7 +641,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -713,7 +660,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -727,17 +673,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ] + "tensor" ], "shape": [ 1024 @@ -754,7 +696,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -768,7 +709,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -777,7 +717,6 @@ "partition_spec": [ [ "fsdp", - "tensor_transpose", "context", "expert" ], @@ -806,7 +745,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -820,7 +758,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -829,7 +766,6 @@ "partition_spec": [ [ "fsdp", - "tensor_transpose", "context", "expert" ], @@ -858,7 +794,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -872,7 +807,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -888,7 +822,6 @@ "stage", [ "fsdp", - "tensor_transpose", "context", "expert" ] @@ -910,7 +843,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -924,17 +856,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ], + "tensor", "stage" ], "shape": [ @@ -953,7 +881,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -967,17 +894,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ], + "tensor", "stage" ], "shape": [ @@ -996,7 +919,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1010,7 +932,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -1026,7 +947,6 @@ "stage", [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -1050,7 +970,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1064,17 +983,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ], + "tensor", "stage" ], "shape": [ @@ -1093,7 +1008,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1107,7 +1021,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -1116,7 +1029,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -1147,7 +1059,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1161,7 +1072,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -1177,7 +1087,6 @@ "stage", [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -1201,7 +1110,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1215,17 +1123,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ], + "tensor", "stage" ], "shape": [ @@ -1244,7 +1148,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1258,7 +1161,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -1274,7 +1176,6 @@ "stage", [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -1298,7 +1199,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1312,7 +1212,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -1321,7 +1220,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -1348,7 +1246,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1362,17 +1259,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ] + "tensor" ], "shape": [ 1024 @@ -1389,7 +1282,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1403,7 +1295,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -1412,7 +1303,6 @@ "partition_spec": [ [ "fsdp", - "tensor_transpose", "context", "expert" ], @@ -1441,7 +1331,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1455,7 +1344,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -1464,7 +1352,6 @@ "partition_spec": [ [ "fsdp", - "tensor_transpose", "context", "expert" ], @@ -1493,7 +1380,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1507,7 +1393,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -1523,7 +1408,6 @@ "stage", [ "fsdp", - "tensor_transpose", "context", "expert" ] @@ -1545,7 +1429,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1559,17 +1442,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ], + "tensor", "stage" ], "shape": [ @@ -1588,7 +1467,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1602,17 +1480,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ], + "tensor", "stage" ], "shape": [ @@ -1631,7 +1505,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1645,7 +1518,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -1661,7 +1533,6 @@ "stage", [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -1685,7 +1556,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1699,17 +1569,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ], + "tensor", "stage" ], "shape": [ @@ -1728,7 +1594,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1742,7 +1607,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -1751,7 +1615,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -1782,7 +1645,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1796,7 +1658,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -1812,7 +1673,6 @@ "stage", [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -1836,7 +1696,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1850,17 +1709,13 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 } }, "partition_spec": [ - [ - "tensor", - "tensor_transpose" - ], + "tensor", "stage" ], "shape": [ @@ -1879,7 +1734,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1893,7 +1747,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -1909,7 +1762,6 @@ "stage", [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -1933,7 +1785,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1947,7 +1798,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1 @@ -1956,7 +1806,6 @@ "partition_spec": [ [ "tensor", - "tensor_transpose", "tensor_sequence", "autoregressive" ], @@ -1983,7 +1832,6 @@ "context", "context_autoregressive", "tensor", - "tensor_transpose", "tensor_sequence", "expert", "autoregressive" @@ -1997,7 +1845,6 @@ "context": 1, "context_autoregressive": 1, "tensor": 1, - "tensor_transpose": 1, "tensor_sequence": 1, "expert": 1, "autoregressive": 1