python310Packages.jax: 0.3.6 -> 0.3.16

[mcwitt]

Description of changes

Upgrades jax to 0.3.16 and jaxlib to 0.3.15.

Still a draft because this depends on #183052.

Note: This marks the darwin build as broken because jax>=0.3.14 sets macos_minimum_os=10.14, but as far as I can tell 10.12 is the latest supported for x86_64 in nixpkgs (relevant issue: #101229).

Things done

• Built on platform(s)
  • x86_64-linux
  • aarch64-linux
  • x86_64-darwin
  • aarch64-darwin
• For non-Linux: Is sandbox = true set in nix.conf? (See Nix manual)
• Tested, as applicable:
  • NixOS test(s) (look inside nixos/tests)
  • and/or package tests
  • or, for functions and "core" functionality, tests in lib/tests or pkgs/test
  • made sure NixOS tests are linked to the relevant packages
• Tested compilation of all packages that depend on this change using nix-shell -p nixpkgs-review --run "nixpkgs-review rev HEAD". Note: all changes have to be committed, also see nixpkgs-review usage
• Tested basic functionality of all binary files (usually in ./result/bin/)
• 22.11 Release Notes (or backporting 22.05 Release notes)
  • (Package updates) Added a release notes entry if the change is major or breaking
  • (Module updates) Added a release notes entry if the change is significant
  • (Module addition) Added a release notes entry if adding a new NixOS module
  • (Release notes changes) Ran nixos/doc/manual/md-to-db.sh to update generated release notes
• Fits CONTRIBUTING.md.

[mcwitt]

Status update: I'm stuck trying to get this to build on x86_64 darwin, getting errors like the following:

external/org_tensorflow/tensorflow/stream_executor/allocator_stats.h:45:3: error: no template named 'optional' in namespace 'std'; did you mean 'absl::optional'?

I'm less familiar with darwin, but it seems like this might be an issue with nixpkgs using an unsupported darwin SDK version for x86_64. Since v0.3.14, jax has bumped macos_minimum_os to 10.14, but as far as I can tell nixpkgs currently only supports 10.12 for x86_64 (relevant issue: #101229).

For now I'm going to mark this broken on darwin since I haven't been able to figure out a workaround.

[samuela]

Thanks for putting this together @mcwitt! Overall this looks thorough and most of the changes make sense to me. Just a few questions here and there...

[samuela]

I was under the impression that the Darwin build was being disabled until we get the macOS 10.14 SDKs?

[mcwitt]

This came from syncing with the tensorflow derivation (here:

nixpkgs/pkgs/development/python-modules/tensorflow/default.nix

Lines 232 to 234 in 3524da2

	] ++ lib.optionals (!stdenv.isDarwin) [
	nsync
	];

)
My thought was to leave the darwin-specific workarounds in place (even though the build is disabled for now) in case switching to the newer SDK fixes the remaining issues. I'm not really sure of the recommended practice in this situation, though -- happy to remove the gating if that seems preferable (and I guess we may end up doing it anyway if the SDK upgrade also fixes the nsync build on darwin).

[samuela]

Gotcha, yeah now that we will have darwin builds working soon (shout out to @uri-canva!), we can figure how to make this work after all

[samuela]

Does this mean that the package fails to build on darwin entirely, or that it fails to build without this particular line of code? Or does it mean that the vendored version of nsync fails to build?

[mcwitt]

I think it means that nsync in nixpkgs is not supported on darwin (darwin is not listed in platforms:

nixpkgs/pkgs/development/libraries/nsync/default.nix

Lines 25 to 27 in 603db2c

	# On macOS we get an error for some reason:
	# > mkdir: cannot create directory 'build': File exists
	platforms = lib.platforms.linux;

)

(This was taken from the tensorflow derivation.)

[uri-canva]

The darwin build might be ok after #184395, since that updates the package to use SDK 11.0, which is available in nixpkgs, it's just not the default one.

[samuela]

The darwin build might be ok after #184395, since that updates the package to use SDK 11.0, which is available in nixpkgs, it's just not the default one.

Once we get #184395 merged, we should def rebase this on top of that

[mcwitt]

Result of nixpkgs-review pr 183051 run on x86_64-linux 1

I skimmed over the breakages. Many seem to be caused by the following:

• chex: broken by upstream removal of tree_multimap. Fixed by upgrading to latest (0.1.4)
• jmp: broken by upstream removal of tree_multimap. Fixed by upgrading to latest (google-deepmind/jmp@260e5ba)
• distrax: many tests fail; upgrade to latest not possible since it adds a numpy<1.23 version bound
• tensorflow-datasets: long-running test suite; not sure if it'll be fixed by the above fixes or a version bump

The others seem to be broken by unrelated issues (mostly incompatible scipy and rich versions).

[mcwitt]

Unfortunately rebasing with #184395 merged hasn't magically fixed the darwin build.

• Bazel fails to find libtool. I think this is fixed by adding cctools to nativeBuildInputs (mcwitt@4f8786c)
• ERROR: /private/tmp/nix-build-bazel-build-jaxlib-0.3.15.drv-0/output/external/llvm-project/mlir/BUILD.bazel:6549:11: Compiling mlir/lib/TableGen/Argument.cpp failed: undeclared inclusion(s) in rule '@llvm-project//mlir:TableGen' (log). @uri-canva do you have any insight here?

[uri-canva]

I have not seen the undeclared inclusion(s) in rule in my testing. Did you test if jaxlib works on its own without this branch? Maybe it's already broken in master.

[uri-canva]

It shows as successful 9 hours ago: https://hydra.nixos.org/eval/1777705?filter=jaxlib&compare=1777668&full=

[uri-canva]

Oh I see, it's probably some new transitive dependency from the version upgrade.

[uri-canva]

Seems to be a known issue with bazel:
tensorflow/tensorflow#55563
bazelbuild/bazel#15359

[uri-canva]

The workarounds might be a bit hard for us though, not sure how well gcc works on darwin, but we can try.

[samuela]

I skimmed over the breakages. Many seem to be caused by the following:

Thanks for digging into these failures @mcwitt!

• tree_multimap issues (affecting chex, jmp): I'm ok letting these slide for now as long as someone promises to come back and do the necessary upgrades after this PR is merged.
• distrax: appears to be a pre-existing failure on master: https://hydra.nixos.org/build/188060343
• tensorflow-datasets: this is an existing failures on master: https://github.com/samuela/nixpkgs-upkeep/runs/8006479894?check_suite_focus=true

Unfortunately rebasing with #184395 merged hasn't magically fixed the darwin build.

Just to confirm: @mcwitt are you building on x86_64-darwin?

[mcwitt]

Just to confirm: @mcwitt are you building on x86_64-darwin?

@samuela yes, the undeclared inclusion(s) error occurs for x86_64-darwin. Here's the ofborg build log: https://logs.nix.ci/?key=nixos/nixpkgs.183051&attempt_id=fe336da2-c899-425e-b4fc-a0e9328d45b0

[samuela]

It appears that python310Packages.jax.x86_64-darwin is also failing on master but due to a different failure: https://hydra.nixos.org/build/188060892/nixlog/1

[samuela]

Is there any hope trying to fix or patch the x86_64-darwin bug ourselves? It appears to be a quite hairy issue in bazel IIUC... Perhaps we should also raise it with the JAX folks?

[mcwitt]

I'm ok letting these slide for now as long as someone promises to come back and do the necessary upgrades after this PR is merged.

@samuela I'd be happy to handle upgrading chex and jmp after this is merged (already have the updates building in local branches so it shouldn't be too much work 🤞 )

[samuela]

@mcwitt @uri-canva How do ya'll feel about merging this as is, despite the fact that it breaks the x86_64-darwin build? On the one hand I hate to break the build... OTOH a) this breakage isn't our fault, b) I don't want PRs to get stuck forever. wdyt?

[uri-canva]

Yeah I think it's ok, I only just fixed the darwin build, and I'm not sure how long it was broken before that, so I don't think there would be anyone depending on it.

[uri-canva]

Let's mark it as broken and add a comment with all the context / links to the issues for future reference. That way when someone looks at it in the future they won't have to start figuring it out from scratch.

[samuela]

Yeah, I def don't want you to think that the x86_64-darwin work was for nothing @uri-canva! We should def commit to getting darwin support back up and running as soon as upstream progress on this issue allows.

[SuperSandro2000]

so we can merge this now or not?

[samuela]

I think we can!