forked from apple/axlearn
-
Notifications
You must be signed in to change notification settings - Fork 0
/
pyproject.toml
176 lines (162 loc) · 6.33 KB
/
pyproject.toml
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
[build-system]
requires = ["flit_core >=3.2,<4"]
build-backend = "flit_core.buildapi"
[project]
name = "axlearn"
version = "0.0.1"
description = "AXLearn"
readme = "README.md"
requires-python = ">=3.8"
# Production dependencies.
# Every time we upgrade JAX, we should try to bring the rest to the newest versions.
dependencies = [
"attrs>=21.3.0",
"absl-py",
"chex>=0.1.6",
"flax>=0.6.5", # only for checkpoints, FrozenDict
"importlab==0.7", # breaks pytype on 0.8
"jax==0.4.13",
"jaxlib==0.4.13",
"nltk==3.7", # for text preprocessing
"numpy<1.24", # needed to pin to < 1.24; tf ragged_tensor depends on deprecated np.object.
"optax==0.1.3", # optimizers (0.1.0 has known bugs).
"portpicker",
"protobuf==3.20.3", # needed to pin to 3.20.x or lower for tensorflow 2.8.0.
"tensorboard-plugin-profile==2.8.0", # needed to pin to 2.8.0 for tensorflow 2.8.0.
"tensorflow==2.8.0; platform_machine == 'x86_64'", # needed for tfds to work
"tensorflow-datasets==4.7.0",
"tensorflow-io[tensorflow]; platform_machine == 'x86_64'",
"tensorflow-io-gcs-filesystem[tensorflow]; platform_machine == 'x86_64'",
"tensorflow_text==2.8.2", # implied by seqio, but also used directly for text processing
"tensorstore>=0.1.21", # used for supporting GDA checkpoints
"typing-extensions>=4.2.0", # needed for typing.Protocol
"seqio@git+https://github.com/google/seqio.git@8efd5dc016bb626509bdf656e2bbcdfa8a7cf2d8", # used for inputs
]
[project.optional-dependencies]
# Apple Silicon dependencies.
apple-silicon = [
"attrs>=21.3.0",
"absl-py",
"chex>=0.1.6",
"flax>=0.6.5", # only for checkpoints, FrozenDict
"jax==0.4.13",
"jaxlib==0.4.13",
"nltk==3.7", # for text preprocessing
"optax>=0.1.1", # optimizers (0.1.0 has known bugs).
"portpicker",
"tensorboard-plugin-profile==2.8.0", # needed to pin to 2.8.0 for tensorflow 2.8.0.
"tensorflow-datasets==4.7.0",
"tensorstore>=0.1.21", # used for supporting GDA checkpoints
"typing-extensions>=4.2.0", # needed for typing.Protocol
"seqio@git+https://github.com/google/seqio.git@8efd5dc016bb626509bdf656e2bbcdfa8a7cf2d8", # used for inputs
]
# Requirements for testing and development.
dev = [
"black==23.1a1", # formatting
"diffusers==0.16.1", # diffuser dependency for golden config updates
"einops",
"evaluate",
"fairseq==0.12.2", # test-only; note Apple Silicon start.md instructions parse this specifier '==', please be careful when upgrading.
"isort", # formatting
"pre-commit", # local pre commit hooks
"pycocotools", # COCO evaluation tools
"pylint==2.13.9", # formatting. Upgrading to 2.14.0 shows (likely) pylint bug.
"pytest", # test runner
"pytest-xdist", # pytest plugin for test parallelism
"pytype==2022.4.22", # type checking
"scikit-learn>=1.1.1", # test-only
"scipy",
"sentencepiece != 0.1.92",
"tqdm", # test-only
"timm==0.6.12", # DiT Dependency test-only
"torch>=1.12.1", # test-only
"torchvision==0.14.1", # test-only
"transformers==4.27.1", # test-only
"wandb", # test-only
"wrapt", # implied by tensorflow-datasets, but also used in config tests.
]
# Note: Currently tested on x86.
t5x = [
"t5==0.9.4",
"t5x@git+https://github.com/google-research/t5x.git@caa500771af314e3cb877b5c158c1db4ae0427c2"
]
[tool.flit.module]
# This defines the import name. https://flit.pypa.io/en/stable/pyproject_toml.html#module-section
name = "axlearn"
# TODO(markblee): Enable once CLI is migrated.
# [project.scripts]
# Entrypoint for CLI.
# axlearn = "axlearn.cli:main"
[tool.pylint.basic]
# A comma-separated list of package or module names from where C extensions may
# be loaded. Extensions are loading into the active Python interpreter and may
# run arbitrary code.
extension-pkg-whitelist = "numpy,jax"
[tool.pylint.messages_control]
# Disable the message, report, category or checker with the given id(s). You
# can either give multiple identifiers separated by comma (,) or put this
# option multiple times (only on the command line, not in the configuration
# file where it should appear only once).You can also use "--disable=all" to
# disable everything first and then reenable specific checks. For example, if
# you want to run only the similarities checker, you can use "--disable=all
# --enable=similarities". If you want to run only the classes checker, but have
# no Warning level messages displayed, use"--disable=all --enable=classes
# --disable=W"
disable = [
"blacklisted-name",
"duplicate-code",
"fixme",
"invalid-name",
"missing-docstring",
"no-else-return",
"no-member", # Config objects use dynamic members generated by attrs.
"protected-access",
"protected-name",
"redefined-builtin",
"redefined-outer-name",
"too-few-public-methods",
"too-many-arguments",
"too-many-locals",
"unnecessary-lambda",
"unsubscriptable-object", # pylint considers jax.Array unsubscriptable.
]
[tool.pylint.typecheck]
# List of members which are set dynamically and missed by Pylint inference
# system, and so shouldn't trigger E1101 when accessed.
generated-members = "googleapiclient.*,numpy.*,torch.*,sentencepiece.*,jaxlib.xla_extension.*"
[tool.pylint.format]
indent-string = " "
max-line-length = 100
fail-under = 10
# Disable ignore-long-lines for URLs and gs paths.
ignore-long-lines = "<?(http|https|gs)://\\S+>?"
[tool.pylint.similarities]
# Don't include imports or docstrings when looking for duplicate-code.
ignore-imports = "yes"
ignore-docstrings = "yes"
# By default, this is 4, which causes many false-positives.
min-similarity-lines = 20
[tool.pylint.MAIN]
load-plugins = [
"pylint.extensions.docparams"
]
[tool.black]
line-length = 100
target-version = ['py38', 'py39']
[tool.pytest.ini_options]
addopts = "-rs -s -p no:warnings --junitxml=test-results/testing.xml"
markers = [
"gs_login: tests needing GS login.",
"tpu: tests needing access to a TPU device.",
"high_cpu: tests that require a lot of CPU.",
"fp64: tests that require 64-bit floating point precision.",
"golden_config: golden config tests.",
"golden_init: golden init tests.",
"golden_regularizer: golden regularizer scale tests.",
"golden_param_update: golden param update scale tests.",
]
testpaths = "."
junit_family="xunit2"
[tool.isort]
line_length = 100
profile = "black"