{tools}[foss/2020b] jax v0.2.19

[deniskristak]

edit (by @boegel): requires easybuilders/easybuild-easyblocks#2534

[Flamefire]

Need another 2 changes in build.py:

• shell([bazel_path, "shutdown"]) needs to be removed or changed to shell([bazel_path] + args.bazel_startup_options + ["shutdown"])
• in the function check_bazel-version change to use --version. See Use bazel --version to check the bazel version jax-ml/jax#7507

[boegel]

@deniskristak The name of the patch files as they are included in the PR are incorrect.

I fixed that in deniskristak#3, along with a couple of other things, like including checksums.

[boegel]

Test report by @boegel
FAILED
Build succeeded for 2 out of 3 (3 easyconfigs in total)
node3111.skitty.os - Linux centos linux 7.9.2009, x86_64, Intel(R) Xeon(R) Gold 6140 CPU @ 2.30GHz, Python 3.6.8
See https://gist.github.com/8ece804dca9e7ea990fb2dbfa0714095 for a full test report.

edit: hmm, one of the tests failed...

_____________________________________________ XlaBridgeTest.test_timer_tpu_warning _____________________________________________
[gw2] linux -- Python 3.8.6 /software/Python/3.8.6-GCCcore-10.2.0/bin/python3.8

self = <xla_bridge_test.XlaBridgeTest testMethod=test_timer_tpu_warning>, _ = <MagicMock name='make_tpu_client' id='47426097730176'>

    @mock.patch('jax.lib.xla_client.make_tpu_client', side_effect=mock_tpu_client)
    def test_timer_tpu_warning(self, _):
      with warnings.catch_warnings(record=True) as w:
        warnings.simplefilter('always')
        xb.tpu_client_timer_callback(0.01)
>       self.assertLen(w, 1)
E       AssertionError: [] has length of 0, expected 1.

tests/xla_bridge_test.py:72: AssertionError

[boegel]

@boegelbot please test @ generoso
CORE_CNT=16

[boegelbot]

@boegel: Request for testing this PR well received on generoso

PR test command 'EB_PR=13622 EB_ARGS= /apps/slurm/default/bin/sbatch --job-name test_PR_13622 --ntasks="16" ~/boegelbot/eb_from_pr_upload_generoso.sh' executed!

• exit code: 0
• output:

Submitted batch job 18132

Test results coming soon (I hope)...

- notification for comment with ID 898704698 processed

Message to humans: this is just bookkeeping information for me,
it is of no use to you (unless you think I have a bug, which I don't).

[boegel]

Test report by @boegel
SUCCESS
Build succeeded for 3 out of 3 (3 easyconfigs in total)
node3111.skitty.os - Linux centos linux 7.9.2009, x86_64, Intel(R) Xeon(R) Gold 6140 CPU @ 2.30GHz, Python 3.6.8
See https://gist.github.com/5891da8c2282beb78f56f429cf1e49f6 for a full test report.

[boegel]

Test report by @boegel
SUCCESS
Build succeeded for 3 out of 3 (3 easyconfigs in total)
node3525.doduo.os - Linux RHEL 8.2, x86_64, AMD EPYC 7552 48-Core Processor (zen2), Python 3.6.8
See https://gist.github.com/dbe92c04b4452fa1055631a989bc9a0a for a full test report.

[boegel]

Test report by @boegel
SUCCESS
Build succeeded for 3 out of 3 (3 easyconfigs in total)
node2607.swalot.os - Linux centos linux 7.9.2009, x86_64, Intel(R) Xeon(R) CPU E5-2660 v3 @ 2.60GHz (haswell), Python 3.6.8
See https://gist.github.com/ef3d0914628552932371ffaca619f0a2 for a full test report.

[boegelbot]

Test report by @boegelbot
SUCCESS
Build succeeded for 3 out of 3 (3 easyconfigs in total)
generoso-x-1 - Linux centos linux 8.2.2004, x86_64, Intel(R) Xeon(R) CPU E5-2667 v3 @ 3.20GHz (haswell), Python 3.6.8
See https://gist.github.com/655945a4cfadc9314f2d179b8f9ed876 for a full test report.

[boegel]

This can/should probably be cleaned up by implementing a custom easyblock for jaxlib.

I'll merge this as is though, since it's working, and a major hurdle towards AlphaFold, and keep that as an exercise for an upcoming PR...

[boegel]

lgtm

[boegel]

Going in, thanks @deniskristak!

[Flamefire]

Test report by @Flamefire
FAILED
Build succeeded for 2 out of 3 (3 easyconfigs in total)
taurusi8002 - Linux centos linux 7.9.2009, x86_64, AMD EPYC 7352 24-Core Processor (zen2), Python 2.7.5
See https://gist.github.com/de777494f99609e7c0c32a739adaf447 for a full test report.

[Flamefire]

Test report by @Flamefire
FAILED
Build succeeded for 2 out of 3 (3 easyconfigs in total)
taurusi6606.taurus.hrsk.tu-dresden.de - Linux RHEL 7.9, x86_64, Intel(R) Xeon(R) CPU E5-2680 v3 @ 2.50GHz (haswell), Python 2.7.5
See https://gist.github.com/13f062a38935f5541d68de71b5113a15 for a full test report.