rocm_jax/build/test-requirements.txt

absl-py
build
cloudpickle
colorama>=0.4.4
filelock
flatbuffers
hypothesis
mpmath>=1.3
pillow>=10.4.0
# TODO(kanglan): Remove once psutil from portpicker supports python 3.13t
portpicker; python_version<"3.13"
pytest-xdist
wheel
rich
setuptools
# matplotlib 3.9.0 pins NumPy 1.23, which is incompatible with the requirement
# below.
matplotlib~=3.8.4; python_version=="3.10"
matplotlib; python_version>="3.11"
opt-einsum
auditwheel
Migrate JAX internals to builtin Python logging This commit changes the JAX codebase to use Python's builtin logging instead of ABSL logging. With the latter being used in JAX code as of now, the change to Python builtin logging is advised for the following reasons (among others): - absl-py can be removed as an external dependency of JAX. - Builtin logging brings the option of adding more log handlers, for example file handlers for log dumps or writers to different IO streams. Logging in JAX is ported over to take place at the module level. While previously, some Python namespaces within JAX already used module-scoped logging via absl.vlog, the following idiom was adopted to provide the same functionality in Python builtin logging: ```py import logging logger = logging.getLogger(__name__) logger.debug(...) logger.info(...) ``` The builtin root logger is left untouched, which is beneficial for downstream users planning to customize the Python root logger. All JAX internal code promises to log to descendants of the top-level "jax" logger by virtue of log propagation. The package `absl-py` was removed from JAX's install requirements, and added into its test requirements. 2022-10-13 17:06:22 +02:00			`absl-py`
Fix warning about direct invocation of setup.py during jaxlib build. The jaxlib wheel build currently uses `python setup.py bdist_wheel` to construct the wheel. Change it to use `python -m build -w` instead. To avoid Python getting confused between the directory named `build` in the bazel tree and the Python `build` module, move `build_wheel.py` into `jaxlib/tools`. PiperOrigin-RevId: 548133811 2023-07-14 08:30:41 -07:00			`build`
Add cloudpickle as a test requirement. We have at least one test that tests pickling JAX objects. 2022-02-16 15:04:56 -05:00			`cloudpickle`
Adds a Wadler-Lindig pretty printer. Changes jaxpr printing to use it. 2021-09-24 22:08:42 -04:00			`colorama>=0.4.4`
Added filelock to test-requirements.txt and requirements lock files This is a follow up to #21741. 2024-06-11 11:53:10 +01:00			`filelock`
[export] Add support for serialization and deserialization of Exported At the moment we can export a JAX function into an Exported and we can invoke an Exported from another JAX function, but there is no way to serialize an Exported to be able to use it in another process. Exported now has all the features we had in mind, so it is a reasonable time to add a serialization method. The intention is for the serialization to have backwards compatibility guarantees, meaning that we can deserialize an Exported that has been serialized with an older version of JAX. This PR does not add explicit support for versioning, nor backwards compatibility tests. Those will follow. Here we add serialization and deserialization to bytearray, using the flatbuffers package. We use flatbuffers because it is simple, it has backwards and forwards compatibility guarantees, it is a lightweight dependency that does not require additional build steps, and it is fast (deserialization simply indexes into the bytearray rather than creating a Python structure). In the process of implementing this we have done some small cleanup of the Exported structure: * renamed serialization_version to mlir_module_serialization_version * renamed disabled_checks to disabled_safety_checks This code is tested by changing export_test.py to interpose a serialization followed by a deserialization every time we export.export. There is a known bug with the serialization of effects, so I disabled one of the export tests. Will fix in a subsequent PR. PiperOrigin-RevId: 590078785 2023-12-11 23:22:16 -08:00			`flatbuffers`
[Pallas] Upstream pallas to JAX PiperOrigin-RevId: 552963029 2023-08-01 16:42:26 -07:00			`hypothesis`
Evaluate the correctness of JAX complex functions using mpmath as a reference 2024-03-21 23:35:29 +02:00			`mpmath>=1.3`
Update pillow version in JAX build test-requirements.txt 2024-09-25 14:55:46 +00:00			`pillow>=10.4.0`
Install test dependencies from test-requirements.txt instead of requirements.in PiperOrigin-RevId: 736878834 2025-03-14 08:56:27 -07:00			`# TODO(kanglan): Remove once psutil from portpicker supports python 3.13t`
Remove portpicker for free threaded python 3.13t in test-requirements.txt PiperOrigin-RevId: 722776783 2025-02-03 13:29:16 -08:00			`portpicker; python_version<"3.13"`
Switch CI builds from Travis to Github actions (#3409) 2020-06-11 17:10:56 -04:00			`pytest-xdist`
Attach source info to Jaxpr equations. (#3421) * Attach source info to Jaxpr equations. Example: ``` In [1]: import jax, jax.numpy as jnp In [2]: def f(x, y): ...: z = jax.numpy.cos(x) ...: z = z * jax.numpy.tanh(y) ...: return z + 2 ...: In [3]: jax.make_jaxpr(jax.value_and_grad(f))(7., 9.) Out[3]: { lambda ; a b. let c = cos a [<ipython-input-2-5d59f71cb65d>:2 (f)] d = tanh b [<ipython-input-2-5d59f71cb65d>:3 (f)] e = mul c d [<ipython-input-2-5d59f71cb65d>:3 (f)] f = add e 2.0 [<ipython-input-2-5d59f71cb65d>:4 (f)] g = mul 1.0 d [<ipython-input-2-5d59f71cb65d>:3 (f)] h = neg g [<ipython-input-2-5d59f71cb65d>:2 (f)] i = sin a [<ipython-input-2-5d59f71cb65d>:2 (f)] j = mul h i [<ipython-input-2-5d59f71cb65d>:2 (f)] in (f, j) } In [7]: print(jax.xla_computation(jax.value_and_grad(f))(7., 9.).as_hlo_module().to_string()) HloModule xla_computation_f__4.15 ENTRY %xla_computation_f__4.15 (parameter.1: f32[], parameter.2: f32[]) -> (f32[], f32[]) { %constant.3 = pred[] constant(false) %parameter.1 = f32[] parameter(0) %cosine.4 = f32[] cosine(f32[] %parameter.1), metadata={op_type="cos" op_name="xla_computation(f)/cos" source_file="<ipython-input-2-5d59f71cb65d>" source_line=2} %parameter.2 = f32[] parameter(1) %tanh.5 = f32[] tanh(f32[] %parameter.2), metadata={op_type="tanh" op_name="xla_computation(f)/tanh" source_file="<ipython-input-2-5d59f71cb65d>" source_line=3} %multiply.6 = f32[] multiply(f32[] %cosine.4, f32[] %tanh.5), metadata={op_type="mul" op_name="xla_computation(f)/mul" source_file="<ipython-input-2-5d59f71cb65d>" source_line=3} %constant.7 = f32[] constant(2), metadata={op_type="add" op_name="xla_computation(f)/add" source_file="<ipython-input-2-5d59f71cb65d>" source_line=4} %add.8 = f32[] add(f32[] %multiply.6, f32[] %constant.7), metadata={op_type="add" op_name="xla_computation(f)/add" source_file="<ipython-input-2-5d59f71cb65d>" source_line=4} %constant.9 = f32[] constant(1), metadata={op_type="mul" op_name="xla_computation(f)/mul" source_file="<ipython-input-2-5d59f71cb65d>" source_line=3} %multiply.10 = f32[] multiply(f32[] %constant.9, f32[] %tanh.5), metadata={op_type="mul" op_name="xla_computation(f)/mul" source_file="<ipython-input-2-5d59f71cb65d>" source_line=3} %negate.11 = f32[] negate(f32[] %multiply.10), metadata={op_type="neg" op_name="xla_computation(f)/neg" source_file="<ipython-input-2-5d59f71cb65d>" source_line=2} %sine.12 = f32[] sine(f32[] %parameter.1), metadata={op_type="sin" op_name="xla_computation(f)/sin" source_file="<ipython-input-2-5d59f71cb65d>" source_line=2} %multiply.13 = f32[] multiply(f32[] %negate.11, f32[] %sine.12), metadata={op_type="mul" op_name="xla_computation(f)/mul" source_file="<ipython-input-2-5d59f71cb65d>" source_line=2} ROOT %tuple.14 = (f32[], f32[]) tuple(f32[] %add.8, f32[] %multiply.13) } ``` Co-authored-by: Matthew Johnson <mattjj@google.com> 2020-06-17 19:35:36 -04:00			`wheel`
Add function to visualize `Sharding`s 2022-09-15 19:04:49 -07:00			`rich`
Remove version limit for `setuptools` dependency. PiperOrigin-RevId: 735453796 2025-03-10 11:34:16 -07:00			`setuptools`
Reverts 72f9a493589a1046e6927a5f16d7dc71df530743 PiperOrigin-RevId: 691843537 2024-10-31 10:04:38 -07:00			`# matplotlib 3.9.0 pins NumPy 1.23, which is incompatible with the requirement`
			`# below.`
Set minimum supported Python version to 3.10 for matplotlib. Temporary fixes an issue with `python -m build` that fails when python 3.8 is used because `matplotlib~=3.8.4` is unavailable for this python version. We are working on creating Bazel build rule with the hermetic Python for JAX wheel ([we already have Jaxlib and plugins build rules ready](https://github.com/jax-ml/jax/pull/23276)). The required python modules are provided in requirements.in file, so when we implement Bazel build rule for JAX wheel, requirements.in will be the only source of dependencies, and test-requirements.txt won't be needed for building JAX wheel. PiperOrigin-RevId: 692260046 2024-11-01 12:33:58 -07:00			`matplotlib~=3.8.4; python_version=="3.10"`
Reverts 72f9a493589a1046e6927a5f16d7dc71df530743 PiperOrigin-RevId: 691843537 2024-10-31 10:04:38 -07:00			`matplotlib; python_version>="3.11"`
			`opt-einsum`
Pin numpy~=2.1.0 in workflow file instead of test-requirements.txt PiperOrigin-RevId: 737632771 2025-03-17 08:58:18 -07:00			`auditwheel`