python311Packages.distrax: fix type errors

[stephen-huan]

Description of changes

As mentioned in #297146 (comment), an update to chex breaks some of distrax's tests.

This PR is a stopgap measure to fix the build until this PR to distrax lands.

Things done

• Built on platform(s)
  • x86_64-linux
  • aarch64-linux
  • x86_64-darwin
  • aarch64-darwin
• For non-Linux: Is sandboxing enabled in nix.conf? (See Nix manual)
  • sandbox = relaxed
  • sandbox = true
• 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/)
• 24.05 Release Notes (or backporting 23.05 and 23.11 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
• Fits CONTRIBUTING.md.

---

Add a 👍 reaction to pull requests you find important.

[GaetanLepage]

nixpkgs-review result

Generated using nixpkgs-review.

Command: nixpkgs-review pr 300648

---

x86_64-linux

✅ 4 packages built:

• python311Packages.distrax
• python311Packages.distrax.dist
• python311Packages.rlax
• python311Packages.rlax.dist

---

aarch64-linux

✅ 4 packages built:

• python311Packages.distrax
• python311Packages.distrax.dist
• python311Packages.rlax
• python311Packages.rlax.dist

---

x86_64-darwin

---

aarch64-darwin

⏩ 4 packages marked as broken and skipped:

• python311Packages.distrax
• python311Packages.distrax.dist
• python311Packages.rlax
• python311Packages.rlax.dist