Skip to content
This repository was archived by the owner on Oct 31, 2025. It is now read-only.
Open
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
109 commits
Select commit Hold shift + click to select a range
d2a6fd9
+ compatibility with CUDA 11.0
mmarcinmichal Jun 30, 2023
48f097c
+ compatibility with CUDA 11.0
mmarcinmichal Jun 30, 2023
8cc0b42
Update build.yaml
mmarcinmichal Aug 28, 2023
af6332b
Update build.yaml
mmarcinmichal Aug 28, 2023
aa354f0
Update build.yaml
mmarcinmichal Aug 28, 2023
63d7990
Update build.yaml
mmarcinmichal Aug 28, 2023
a72422d
Update rnn_test.py
mmarcinmichal Aug 29, 2023
605f122
Update rnn_test.py
mmarcinmichal Aug 29, 2023
b9d5739
+ add test_fork.py
mmarcinmichal Sep 28, 2023
e82f33b
Merge branch 'master' of https://github.com/mmarcinmichal/trax
mmarcinmichal Sep 28, 2023
cc46a6d
+ build
mmarcinmichal Oct 2, 2023
34abf50
+ build
mmarcinmichal Oct 2, 2023
de9862f
+ build
mmarcinmichal Oct 2, 2023
02b8d17
+ build
mmarcinmichal Oct 2, 2023
f53b338
+ build
mmarcinmichal Oct 2, 2023
9c849bf
+ build
mmarcinmichal Oct 2, 2023
cac2053
+ build
mmarcinmichal Oct 2, 2023
698d0fd
+ build
mmarcinmichal Oct 2, 2023
241ce7e
+ build
mmarcinmichal Oct 2, 2023
ad1d76d
- Refactoring combatible with the following runtime enviroment
mmarcinmichal Oct 4, 2023
918f8a2
- Refactoring combatible with the following runtime enviroment
mmarcinmichal Oct 6, 2023
2b62eb8
- Refactoring combatible with the following runtime enviroment
mmarcinmichal Oct 16, 2023
244b522
- Refactoring combatible with the following runtime enviroment
mmarcinmichal Oct 16, 2023
7589fd2
- Refactoring combatible with the following runtime enviroment
mmarcinmichal Oct 16, 2023
d8a5605
- Refactoring combatible with the following runtime enviroment
mmarcinmichal Oct 16, 2023
31c060c
- Refactoring combatible with the following runtime enviroment
mmarcinmichal Oct 16, 2023
59abb4d
- Refactoring combatible with the following runtime enviroment
mmarcinmichal Oct 16, 2023
3aa35cf
- Refactoring combatible with the following runtime enviroment
mmarcinmichal Oct 16, 2023
1d57eaf
- Refactoring combatible with the following runtime enviroment
mmarcinmichal Oct 16, 2023
a92098b
- Refactoring combatible with the following runtime enviroment
mmarcinmichal Oct 16, 2023
fb0302f
- Refactoring combatible with the following runtime enviroment
mmarcinmichal Oct 16, 2023
fd6abf6
- Refactoring combatible with the following runtime enviroment
mmarcinmichal Oct 16, 2023
b3d2871
- Refactoring combatible with the following runtime enviroment
mmarcinmichal Oct 16, 2023
6f35bab
- Refactoring combatible with the following runtime enviroment
mmarcinmichal Oct 16, 2023
6b12922
- Refactoring combatible with the following runtime enviroment
mmarcinmichal Oct 16, 2023
b5853d0
- Refactoring combatible with the following runtime enviroment
mmarcinmichal Oct 16, 2023
338f526
- Refactoring combatible with the following runtime enviroment
mmarcinmichal Oct 16, 2023
62dda26
- Refactoring combatible with the following runtime enviroment
mmarcinmichal Oct 16, 2023
baf28b8
- Refactoring combatible with the following runtime enviroment
mmarcinmichal Oct 16, 2023
a1db8c8
- Refactoring combatible with the following runtime enviroment
mmarcinmichal Oct 16, 2023
d801be5
- Refactoring combatible with the following runtime enviroment
mmarcinmichal Oct 16, 2023
f9f6189
- Refactoring combatible with the following runtime enviroment
mmarcinmichal Oct 16, 2023
a2921a1
- Refactoring combatible with the following runtime enviroment
mmarcinmichal Oct 16, 2023
7fce6a3
- Refactoring combatible with the following runtime enviroment
mmarcinmichal Oct 16, 2023
e0f5de1
- Refactoring combatible with the following runtime enviroment
mmarcinmichal Oct 17, 2023
b524feb
- Refactoring combatible with the following runtime enviroment
mmarcinmichal Oct 17, 2023
9678b54
- Refactoring combatible with the following runtime enviroment
mmarcinmichal Oct 17, 2023
0d13004
- Refactoring combatible with the following runtime enviroment
mmarcinmichal Oct 17, 2023
d08d535
- Refactoring combatible with the following runtime enviroment
mmarcinmichal Oct 17, 2023
0c6a7ed
- Refactoring combatible with the following runtime enviroment
mmarcinmichal Oct 17, 2023
a1c0168
- Build information
mmarcinmichal Oct 19, 2023
e4582d1
- Build information
mmarcinmichal Oct 20, 2023
3fcd8d5
- Build information
mmarcinmichal Oct 20, 2023
c505bfb
- New version to release
mmarcinmichal Nov 14, 2023
90e42cc
Create codeql.yml
mmarcinmichal Nov 14, 2023
11ac771
- TF 2.14.0
mmarcinmichal Nov 24, 2023
18b0c3d
- TF 2.14.0
mmarcinmichal Nov 24, 2023
5d6b823
- TF 2.14.0
mmarcinmichal Nov 24, 2023
7aad754
- TF 2.14.0
mmarcinmichal Nov 24, 2023
80d7bc1
- TF 2.14.0
mmarcinmichal Nov 24, 2023
fbbaf43
- TF 2.14.0
mmarcinmichal Nov 24, 2023
2ddfaf2
- TF 2.11.0
mmarcinmichal Nov 24, 2023
eb06512
- TF 2.11.0
mmarcinmichal Nov 24, 2023
8819636
- TF 2.13.0
mmarcinmichal Nov 24, 2023
84291ff
- numpy==1.23.5
mmarcinmichal Nov 27, 2023
497b570
- numpy==1.23.5
mmarcinmichal Nov 27, 2023
73a3cf5
- numpy==1.23.5
mmarcinmichal Dec 1, 2023
cf412bb
Update build.yaml
mmarcinmichal Feb 24, 2025
122150e
Update build.yaml
mmarcinmichal Feb 24, 2025
a1e1462
Update build.yaml
mmarcinmichal Feb 24, 2025
592bca9
Update build.yaml
mmarcinmichal Feb 24, 2025
f5a38ad
Update build.yaml
mmarcinmichal Feb 24, 2025
5b4d43f
Update build.yaml
mmarcinmichal Feb 24, 2025
22d0377
Update setup.py
mmarcinmichal Mar 25, 2025
9056bdc
Update setup.py
mmarcinmichal Mar 25, 2025
2219290
+ working bert model i multi devices gpus
mmarcinmichal Oct 19, 2023
54a76ac
Refactoring of data module + test
mmarcinmichal Apr 2, 2025
ee632b0
Refactoring of data, fastmath, tf modules + tests
mmarcinmichal Apr 2, 2025
ed0409c
Refactoring of full layer package test - all working fine
mmarcinmichal Apr 3, 2025
fcb7c9b
Refactoring of full model package test - all working fine
mmarcinmichal Apr 3, 2025
ac99b05
Refactoring of full package. The most important things is switch trai…
mmarcinmichal Apr 8, 2025
53d077c
Refactoring of full package. The most important things is switch trai…
mmarcinmichal Apr 8, 2025
f44a442
Refactoring of full package. The most important things is switch trai…
mmarcinmichal Apr 8, 2025
9f6b334
Refactoring of full package. Reformmer and terafformer tests are passed
mmarcinmichal Apr 8, 2025
80436a1
Refactoring of full package. Reformmer and terafformer tests are passed
mmarcinmichal Apr 9, 2025
cc08a4a
Refactoring of full package. Reformmer and terafformer tests are passed
mmarcinmichal Apr 9, 2025
31aaef5
Refactoring of full package. Reformmer and terafformer tests are passed
mmarcinmichal Apr 9, 2025
8186233
Refactoring of full package. Reformmer and terafformer tests are passed
mmarcinmichal Apr 10, 2025
7335ef5
Yaml build refactoriung
mmarcinmichal Apr 10, 2025
5351746
Yaml build refactoriung
mmarcinmichal Apr 10, 2025
e19f28b
Yaml build refactoriung
mmarcinmichal Apr 10, 2025
3fb18fe
Yaml build refactoriung
mmarcinmichal Apr 10, 2025
1eda928
Add Missing seqio
mmarcinmichal Apr 10, 2025
0b42d2f
Test improvementd
mmarcinmichal Apr 10, 2025
a90b00b
Clear outputs
mmarcinmichal Apr 11, 2025
d060300
Clear outputs
mmarcinmichal Apr 11, 2025
fa01e11
Clear outputs
mmarcinmichal Apr 11, 2025
5aadab6
Clear outputs
mmarcinmichal Apr 11, 2025
dd53be2
Examples refactoring
mmarcinmichal Apr 28, 2025
888c54c
Add flash attention implementation and tests
mmarcinmichal Jun 19, 2025
938a860
Add GNN classification examples
mmarcinmichal Jun 23, 2025
63a3629
Examples refactoring
mmarcinmichal Aug 26, 2025
a7b64db
Improve GNN 20NG training
mmarcinmichal Aug 27, 2025
70643d4
Examples refactoring
mmarcinmichal Nov 4, 2025
1fb9e09
Update build.yaml
mmarcinmichal Nov 5, 2025
ca0f7b2
Update runner version to Ubuntu 22.04
mmarcinmichal Nov 6, 2025
c25e455
Update build.yaml
mmarcinmichal Nov 6, 2025
1910e39
ubuntu-22.04 and ligweights imports
mmarcinmichal Nov 7, 2025
997ff52
ubuntu-22.04 and ligweights imports
mmarcinmichal Nov 7, 2025
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
The table of contents is too big for display.
Diff view
Diff view
  •  
  •  
  •  
100 changes: 49 additions & 51 deletions .github/workflows/build.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -21,9 +21,9 @@ name: build
on:
# Triggers the workflow on push or pull request events but only for the master branch
push:
branches: [ master ]
branches: [ "master" ]
pull_request:
branches: [ master ]
branches: [ "master" ]

# Allows you to run this workflow manually from the Actions tab
workflow_dispatch:
Expand All @@ -33,59 +33,57 @@ jobs:
# This workflow contains a single job called "build"
build:
# The type of runner that the job will run on
runs-on: ubuntu-latest
runs-on: ubuntu-22.04

strategy:
matrix:
python-version: [3.7]
# tf-nightly has some pip version conflicts, so can't be installed.
# Use only numbered TF as of now.
# tf-version: ["2.4.*", "tf-nightly"]
tf-version: ["2.4.*"]
python-version: [ '3.10' ]
# Which tf-version run.
tf-version: [ '2.17.0' ]
# Which set of tests to run.
trax-test: ["lib", "research"]
trax-test: [ 'lib','research' ]

# Steps represent a sequence of tasks that will be executed as part of the job
steps:
# Checks-out your repository under $GITHUB_WORKSPACE, so your job can access it
- uses: actions/checkout@v2
- name: Set up Python ${{ matrix.python-version }}
uses: actions/setup-python@v2
with:
python-version: ${{ matrix.python-version }}
- name: Install dependencies
run: |
python -m pip install --upgrade pip
python -m pip install -q -U setuptools numpy
python -m pip install flake8 pytest
if [[ ${{matrix.tf-version}} == "tf-nightly" ]]; then python -m pip install tf-nightly; else python -m pip install -q "tensorflow=="${{matrix.tf-version}}; fi
pip install -e .[tests,t5]
# # Lint with flake8
# - name: Lint with flake8
# run: |
# # stop the build if there are Python syntax errors or undefined names
# flake8 . --count --select=E9,F63,F7,F82 --show-source --statistics
# # exit-zero treats all errors as warnings. The GitHub editor is 127 chars wide
# flake8 . --count --exit-zero --max-complexity=10 --max-line-length=127 --statistics
# Test out right now with only testing one directory.
- name: Test with pytest
run: |
TRAX_TEST=" ${{matrix.trax-test}}" ./oss_scripts/oss_tests.sh
# The below step just reports the success or failure of tests as a "commit status".
# This is needed for copybara integration.
- name: Report success or failure as github status
if: always()
shell: bash
run: |
status="${{ job.status }}"
lowercase_status=$(echo $status | tr '[:upper:]' '[:lower:]')
curl -sS --request POST \
--url https://api.github.com/repos/${{ github.repository }}/statuses/${{ github.sha }} \
--header 'authorization: Bearer ${{ secrets.GITHUB_TOKEN }}' \
--header 'content-type: application/json' \
--data '{
"state": "'$lowercase_status'",
"target_url": "https://github.com/${{ github.repository }}/actions/runs/${{ github.run_id }}",
"description": "'$status'",
"context": "github-actions/build"
}'
# Checks-out your repository under $GITHUB_WORKSPACE, so your job can access it
- uses: actions/checkout@v3
- name: Set up Python ${{matrix.python-version}}
uses: actions/setup-python@v5
with:
python-version: ${{matrix.python-version}}
cache: 'pip'
- name: Install dependencies
env:
PIP_DISABLE_PIP_VERSION_CHECK: '1'
run: |
python -m pip install -U pip
# Install TensorFlow matching matrix version.
python -m pip install "tensorflow==${{ matrix.tf-version }}"
# Install package in editable mode with test and T5 extras (tests use T5 preprocessors).
python -m pip install -e .[tests,t5,rl]
# Test out right now with only testing one directory.
- name: Install trax package
run: |
python -m pip install -e .
- name: Test with pytest
working-directory: .
run: |
TRAX_TEST="${{matrix.trax-test}}" ./oss_scripts/oss_tests.sh
# The below step just reports the success or failure of tests as a "commit status".
# This is needed for copy bara integration.
- name: Report success or failure as github status
if: always()
shell: bash
run: |
status="${{ job.status }}"
lowercase_status=$(echo $status | tr '[:upper:]' '[:lower:]')
curl -sS --request POST \
--url https://api.github.com/repos/${{github.repository}}/statuses/${{github.sha}} \
--header 'authorization: Bearer ${{ secrets.GITHUB_TOKEN }}' \
--header 'content-type: application/json' \
--data '{
"state": "'$lowercase_status'",
"target_url": "https://github.com/${{github.repository}}/actions/runs/${{github.run_id}}",
"description": "'$status'",
"context": "github-actions/build"
}'
82 changes: 82 additions & 0 deletions .github/workflows/codeql.yml
Original file line number Diff line number Diff line change
@@ -0,0 +1,82 @@
# For most projects, this workflow file will not need changing; you simply need
# to commit it to your repository.
#
# You may wish to alter this file to override the set of languages analyzed,
# or to provide custom queries or build logic.
#
# ******** NOTE ********
# We have attempted to detect the languages in your repository. Please check
# the `language` matrix defined below to confirm you have the correct set of
# supported CodeQL languages.
#
name: "CodeQL"

on:
push:
branches: [ "1.5.1" ]
pull_request:
# The branches below must be a subset of the branches above
branches: [ "1.5.1" ]
schedule:
- cron: '31 4 * * 1'

jobs:
analyze:
name: Analyze
# Runner size impacts CodeQL analysis time. To learn more, please see:
# - https://gh.io/recommended-hardware-resources-for-running-codeql
# - https://gh.io/supported-runners-and-hardware-resources
# - https://gh.io/using-larger-runners
# Consider using larger runners for possible analysis time improvements.
runs-on: ${{ (matrix.language == 'swift' && 'macos-latest') || 'ubuntu-latest' }}
timeout-minutes: ${{ (matrix.language == 'swift' && 120) || 360 }}
permissions:
actions: read
contents: read
security-events: write

strategy:
fail-fast: false
matrix:
language: [ 'python' ]
# CodeQL supports [ 'c-cpp', 'csharp', 'go', 'java-kotlin', 'javascript-typescript', 'python', 'ruby', 'swift' ]
# Use only 'java-kotlin' to analyze code written in Java, Kotlin or both
# Use only 'javascript-typescript' to analyze code written in JavaScript, TypeScript or both
# Learn more about CodeQL language support at https://aka.ms/codeql-docs/language-support

steps:
- name: Checkout repository
uses: actions/checkout@v3

# Initializes the CodeQL tools for scanning.
- name: Initialize CodeQL
uses: github/codeql-action/init@v2
with:
languages: ${{ matrix.language }}
# If you wish to specify custom queries, you can do so here or in a config file.
# By default, queries listed here will override any specified in a config file.
# Prefix the list here with "+" to use these queries and those in the config file.

# For more details on CodeQL's query packs, refer to: https://docs.github.com/en/code-security/code-scanning/automatically-scanning-your-code-for-vulnerabilities-and-errors/configuring-code-scanning#using-queries-in-ql-packs
# queries: security-extended,security-and-quality


# Autobuild attempts to build any compiled languages (C/C++, C#, Go, Java, or Swift).
# If this step fails, then you should remove it and run the build manually (see below)
- name: Autobuild
uses: github/codeql-action/autobuild@v2

# ℹ️ Command-line programs to run using the OS shell.
# 📚 See https://docs.github.com/en/actions/using-workflows/workflow-syntax-for-github-actions#jobsjob_idstepsrun

# If the Autobuild fails above, remove it and uncomment the following three lines.
# modify them (or add more) to build your code if your project, please refer to the EXAMPLE below for guidance.

# - run: |
# echo "Run, Build Application using script"
# ./location_of_script_within_repo/buildscript.sh

- name: Perform CodeQL Analysis
uses: github/codeql-action/analyze@v2
with:
category: "/language:${{matrix.language}}"
2 changes: 1 addition & 1 deletion .readthedocs.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,6 @@ formats: all

# Optionally set the version of Python and requirements required to build your docs
python:
version: 3.7
version: 3.10
install:
- requirements: docs/requirements.txt
4 changes: 2 additions & 2 deletions .travis.yml
Original file line number Diff line number Diff line change
Expand Up @@ -5,10 +5,10 @@ git:
depth: 3
quiet: true
python:
- "3.6"
- "3.10"
env:
global:
- TF_VERSION="2.4.*"
- TF_VERSION="2.11.0"
matrix:
- TRAX_TEST="lib"
- TRAX_TEST="research"
Expand Down
11 changes: 10 additions & 1 deletion README.md
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@
version](https://badge.fury.io/py/trax.svg)](https://badge.fury.io/py/trax)
[![GitHub
Issues](https://img.shields.io/github/issues/google/trax.svg)](https://github.com/google/trax/issues)
![GitHub Build](https://github.com/google/trax/actions/workflows/build.yaml/badge.svg)
![GitHub Build](https://github.com/mmarcinmichal/trax/actions/workflows/build.yaml/badge.svg)
[![Contributions
welcome](https://img.shields.io/badge/contributions-welcome-brightgreen.svg)](CONTRIBUTING.md)
[![License](https://img.shields.io/badge/License-Apache%202.0-brightgreen.svg)](https://opensource.org/licenses/Apache-2.0)
Expand All @@ -26,6 +26,15 @@ Here are a few example notebooks:-
* [**trax.data API explained**](https://github.com/google/trax/blob/master/trax/examples/trax_data_Explained.ipynb) : Explains some of the major functions in the `trax.data` API
* [**Named Entity Recognition using Reformer**](https://github.com/google/trax/blob/master/trax/examples/NER_using_Reformer.ipynb) : Uses a [Kaggle dataset](https://www.kaggle.com/abhinavwalia95/entity-annotated-corpus) for implementing Named Entity Recognition using the [Reformer](https://arxiv.org/abs/2001.04451) architecture.
* [**Deep N-Gram models**](https://github.com/google/trax/blob/master/trax/examples/Deep_N_Gram_Models.ipynb) : Implementation of deep n-gram models trained on Shakespeares works
* **Graph neural networks**: baseline models available via
`trax.models.GraphConvNet`,
`trax.models.GraphEdgeNet` for node and edge updates, or the
attention-based `trax.models.GraphAttentionNet`.
* Example Python scripts using these GNNs for MNIST and IMDB classification are
in
[`resources/examples/python/gnn_mnist/train.py`](resources/examples/python/gnn_mnist/train.py)
and
[`resources/examples/python/gnn_imdb/train.py`](resources/examples/python/gnn_imdb/train.py).



Expand Down
2 changes: 1 addition & 1 deletion docs/.readthedocs.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,6 @@ formats: all

# Optionally set the version of Python and requirements required to build your docs
python:
version: 3.7
version: 3.10
install:
- requirements: docs/requirements.txt
Loading