Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Adding Jit functionality to Plan finialization. #1

Open
wants to merge 34 commits into
base: m1-jax
Choose a base branch
from
Open

Adding Jit functionality to Plan finialization. #1

wants to merge 34 commits into from

Commits on Jul 31, 2024

  1. Adding Jax tests for the M1 mac.

    @alxmrs
    alxmrs committed Jul 31, 2024
    Configuration menu
    Copy the full SHA
    69a906c View commit details
    Browse the repository at this point in the history

  2. Providing flag to default to single precision instead of double for n… 

    …umpy create functions.

This allows us to match the same behavior for Jax on Mac (metal). It seems like this is a desired trait for GPU programming.
    @alxmrs
    alxmrs committed Jul 31, 2024
    Configuration menu
    Copy the full SHA
    d571970 View commit details
    Browse the repository at this point in the history

  3. Matmul on jax-metal only supports floats and complex types. :/

    @alxmrs
    alxmrs committed Jul 31, 2024
    Configuration menu
    Copy the full SHA
    9a043da View commit details
    Browse the repository at this point in the history

  4. Fixes for more tests.

    @alxmrs
    alxmrs committed Jul 31, 2024
    Configuration menu
    Copy the full SHA
    f2afa77 View commit details
    Browse the repository at this point in the history

  5. Run jax tests on PRs.

    @alxmrs
    alxmrs committed Jul 31, 2024
    Configuration menu
    Copy the full SHA
    c0c7751 View commit details
    Browse the repository at this point in the history

  6. Implemented feedback from Tom.

    @alxmrs
    alxmrs committed Jul 31, 2024
    Configuration menu
    Copy the full SHA
    eb059d9 View commit details
    Browse the repository at this point in the history

  7. Added device argument.

    @alxmrs
    alxmrs committed Jul 31, 2024
    Configuration menu
    Copy the full SHA
    fbb78f5 View commit details
    Browse the repository at this point in the history

  8. Exclude 'indexing' case.

    @alxmrs
    alxmrs committed Jul 31, 2024
    Configuration menu
    Copy the full SHA
    83cdd4f View commit details
    Browse the repository at this point in the history

  9. Removing 'indexing' from default types for now.

    @alxmrs
    alxmrs committed Jul 31, 2024
    Configuration menu
    Copy the full SHA
    ff5e13f View commit details
    Browse the repository at this point in the history

  10. Fixed bag argument in test.

    @alxmrs
    alxmrs committed Jul 31, 2024
    Configuration menu
    Copy the full SHA
    a6fcaf3 View commit details
    Browse the repository at this point in the history

  11. Random should respect dtypes.

    @alxmrs
    alxmrs committed Jul 31, 2024
    Configuration menu
    Copy the full SHA
    aa0fdfc View commit details
    Browse the repository at this point in the history

  12. I think I got the linspace tests passing. Need to check if array equa… 

    …lity tolerances are OK.
    @alxmrs
    alxmrs committed Jul 31, 2024
    Configuration menu
    Copy the full SHA
    ecf2a42 View commit details
    Browse the repository at this point in the history

  13. All jax tests pass locally.

    @alxmrs
    alxmrs committed Jul 31, 2024
    Configuration menu
    Copy the full SHA
    f37a5eb View commit details
    Browse the repository at this point in the history

  14. Cleanup for PR.

    @alxmrs
    alxmrs committed Jul 31, 2024
    Configuration menu
    Copy the full SHA
    3adca6b View commit details
    Browse the repository at this point in the history

  15. Added developer instructions for getting set up on the M1+ processor.

    @alxmrs
    alxmrs committed Jul 31, 2024
    Configuration menu
    Copy the full SHA
    f3ab3b3 View commit details
    Browse the repository at this point in the history

  16. Fix env variable expressions.

    @alxmrs
    alxmrs committed Jul 31, 2024
    Configuration menu
    Copy the full SHA
    8c0384a View commit details
    Browse the repository at this point in the history

  17. The macos-14 image requires at least Python 3.10. 

    See actions/setup-python#825.
    @alxmrs
    alxmrs committed Jul 31, 2024
    Configuration menu
    Copy the full SHA
    20cad4b View commit details
    Browse the repository at this point in the history

  18. Looks like 14 specifically supports python 3.11 and up.

    @alxmrs
    alxmrs committed Jul 31, 2024
    Configuration menu
    Copy the full SHA
    bfe0d87 View commit details
    Browse the repository at this point in the history

  19. Splitting install steps for jax for mac and non-mac.

    @alxmrs
    alxmrs committed Jul 31, 2024
    Configuration menu
    Copy the full SHA
    448c60d View commit details
    Browse the repository at this point in the history

  20. Mac install now enables jaxlib compatibility.

    @alxmrs
    alxmrs committed Jul 31, 2024
    Configuration menu
    Copy the full SHA
    b2a16d8 View commit details
    Browse the repository at this point in the history

  21. Put the mac-specific env setting in the right place.

    @alxmrs
    alxmrs committed Jul 31, 2024
    Configuration menu
    Copy the full SHA
    d8c6066 View commit details
    Browse the repository at this point in the history

  22. Test the xlarge runner.

    @alxmrs
    alxmrs committed Jul 31, 2024
    Configuration menu
    Copy the full SHA
    d52fce5 View commit details
    Browse the repository at this point in the history

  23. Extracted to method _to_default_precision.

    @alxmrs
    alxmrs committed Jul 31, 2024
    Configuration menu
    Copy the full SHA
    bdb5777 View commit details
    Browse the repository at this point in the history

  24. Added helpful note.

    @alxmrs
    alxmrs committed Jul 31, 2024
    Configuration menu
    Copy the full SHA
    a9ad5df View commit details
    Browse the repository at this point in the history

  25. No need to cast in the linspace test.

    @alxmrs
    alxmrs committed Jul 31, 2024
    Configuration menu
    Copy the full SHA
    0d436b1 View commit details
    Browse the repository at this point in the history

  26. Remove macos-specific tests.

    @alxmrs
    alxmrs committed Jul 31, 2024
    Configuration menu
    Copy the full SHA
    a67fee2 View commit details
    Browse the repository at this point in the history

  27. Back tracking to previous macos build

    @alxmrs
    alxmrs committed Jul 31, 2024
    Configuration menu
    Copy the full SHA
    2e0a2f9 View commit details
    Browse the repository at this point in the history

  28. Mac should use CPU Jax

    @alxmrs
    alxmrs committed Jul 31, 2024
    Configuration menu
    Copy the full SHA
    a19ca2e View commit details
    Browse the repository at this point in the history

  29. Remove jax mac extra install.

    @alxmrs
    alxmrs committed Jul 31, 2024
    Configuration menu
    Copy the full SHA
    780e349 View commit details
    Browse the repository at this point in the history

  30. Tensordot testing with and without hardcoded types.

    @alxmrs
    alxmrs committed Jul 31, 2024
    Configuration menu
    Copy the full SHA
    2f7c324 View commit details
    Browse the repository at this point in the history

  31. Defaulting to 64 bit precision.

    @alxmrs
    alxmrs committed Jul 31, 2024
    Configuration menu
    Copy the full SHA
    ecf7f15 View commit details
    Browse the repository at this point in the history

  32. rm unneeded test arg.

    @alxmrs
    alxmrs committed Jul 31, 2024
    Configuration menu
    Copy the full SHA
    e189c09 View commit details
    Browse the repository at this point in the history

  33. Revert tensordot change.

    @alxmrs
    alxmrs committed Jul 31, 2024
    Configuration menu
    Copy the full SHA
    9a7a419 View commit details
    Browse the repository at this point in the history

  34. Merged conflicts.

    @alxmrs
    alxmrs committed Jul 31, 2024
    Configuration menu
    Copy the full SHA
    995436b View commit details
    Browse the repository at this point in the history