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

Feature/batched vmap #588

Merged
merged 79 commits into from
Oct 16, 2024
Merged

Feature/batched vmap #588

merged 79 commits into from
Oct 16, 2024

Commits on Oct 10, 2024

  1. Start batched vmap

    @michalk8
    michalk8 committed Oct 10, 2024
    Configuration menu
    Copy the full SHA
    4472425 View commit details
    Browse the repository at this point in the history

  2. Initial batched_vmap impl

    @michalk8
    michalk8 committed Oct 10, 2024
    Configuration menu
    Copy the full SHA
    80b9c60 View commit details
    Browse the repository at this point in the history

  3. Nicer formatting

    @michalk8
    michalk8 committed Oct 10, 2024
    Configuration menu
    Copy the full SHA
    7c58b9a View commit details
    Browse the repository at this point in the history

  4. Fix getting shape

    @michalk8
    michalk8 committed Oct 10, 2024
    Configuration menu
    Copy the full SHA
    b995d15 View commit details
    Browse the repository at this point in the history

  5. Remove private API usage

    @michalk8
    michalk8 committed Oct 10, 2024
    Configuration menu
    Copy the full SHA
    e5cfa1e View commit details
    Browse the repository at this point in the history

  6. Fix new args

    @michalk8
    michalk8 committed Oct 10, 2024
    Configuration menu
    Copy the full SHA
    f8eee7b View commit details
    Browse the repository at this point in the history

  7. Add a TODO

    @michalk8
    michalk8 committed Oct 10, 2024
    Configuration menu
    Copy the full SHA
    51980ef View commit details
    Browse the repository at this point in the history

  8. Canonicalize axes

    @michalk8
    michalk8 committed Oct 10, 2024
    Configuration menu
    Copy the full SHA
    7a510ae View commit details
    Browse the repository at this point in the history

  9. Add batched_vmap to docs

    @michalk8
    michalk8 committed Oct 10, 2024
    Configuration menu
    Copy the full SHA
    1ca05f1 View commit details
    Browse the repository at this point in the history

  10. Removed batched transport functions

    @michalk8
    michalk8 committed Oct 10, 2024
    Configuration menu
    Copy the full SHA
    e7ad8a4 View commit details
    Browse the repository at this point in the history

  11. Remove _norm_{x,y} from CostFn

    @michalk8
    michalk8 committed Oct 10, 2024
    Configuration menu
    Copy the full SHA
    acf221b View commit details
    Browse the repository at this point in the history

  12. Implement apply_lse_kernel

    @michalk8
    michalk8 committed Oct 10, 2024
    Configuration menu
    Copy the full SHA
    07f4734 View commit details
    Browse the repository at this point in the history

  13. Implememt apply_kernel

    @michalk8
    michalk8 committed Oct 10, 2024
    Configuration menu
    Copy the full SHA
    98f1de9 View commit details
    Browse the repository at this point in the history

  14. Implement apply_cost

    @michalk8
    michalk8 committed Oct 10, 2024
    Configuration menu
    Copy the full SHA
    ba825f9 View commit details
    Browse the repository at this point in the history

  15. Remove old functions

    @michalk8
    michalk8 committed Oct 10, 2024
    Configuration menu
    Copy the full SHA
    9036329 View commit details
    Browse the repository at this point in the history

  16. Make function private

    @michalk8
    michalk8 committed Oct 10, 2024
    Configuration menu
    Copy the full SHA
    62089bc View commit details
    Browse the repository at this point in the history

  17. Refactor apply_cost to have consistent shapes

    @michalk8
    michalk8 committed Oct 10, 2024
    Configuration menu
    Copy the full SHA
    b9bb64a View commit details
    Browse the repository at this point in the history

  18. Use _apply_cost_to_vec in PointCloud

    @michalk8
    michalk8 committed Oct 10, 2024
    Configuration menu
    Copy the full SHA
    d8b5ea6 View commit details
    Browse the repository at this point in the history

  19. Remoeve TODO

    @michalk8
    michalk8 committed Oct 10, 2024
    Configuration menu
    Copy the full SHA
    12f923c View commit details
    Browse the repository at this point in the history

  20. Formatting

    @michalk8
    michalk8 committed Oct 10, 2024
    Configuration menu
    Copy the full SHA
    a43dcdb View commit details
    Browse the repository at this point in the history

  21. Simplify _apply_sqeucl_cost

    @michalk8
    michalk8 committed Oct 10, 2024
    Configuration menu
    Copy the full SHA
    e0d75b0 View commit details
    Browse the repository at this point in the history

  22. Fix RecusionError

    @michalk8
    michalk8 committed Oct 10, 2024
    Configuration menu
    Copy the full SHA
    f5445ec View commit details
    Browse the repository at this point in the history

  23. Remove docstring of a private method

    @michalk8
    michalk8 committed Oct 10, 2024
    Configuration menu
    Copy the full SHA
    4922ca9 View commit details
    Browse the repository at this point in the history

  24. Fix apply_lse_kernel

    @michalk8
    michalk8 committed Oct 10, 2024
    Configuration menu
    Copy the full SHA
    799c108 View commit details
    Browse the repository at this point in the history

  25. Squeeze only 1 axis of the cost

    @michalk8
    michalk8 committed Oct 10, 2024
    Configuration menu
    Copy the full SHA
    9a5c1ca View commit details
    Browse the repository at this point in the history

  26. Add TODO

    @michalk8
    michalk8 committed Oct 10, 2024
    Configuration menu
    Copy the full SHA
    8543538 View commit details
    Browse the repository at this point in the history

  27. Rename function, make a property

    @michalk8
    michalk8 committed Oct 10, 2024
    Configuration menu
    Copy the full SHA
    317eb02 View commit details
    Browse the repository at this point in the history

  28. Remove unused helper function

    @michalk8
    michalk8 committed Oct 10, 2024
    Configuration menu
    Copy the full SHA
    d31fd4d View commit details
    Browse the repository at this point in the history

  29. Compute mean summary online

    @michalk8
    michalk8 committed Oct 10, 2024
    Configuration menu
    Copy the full SHA
    4b0f150 View commit details
    Browse the repository at this point in the history

  30. Compute mean online

    @michalk8
    michalk8 committed Oct 10, 2024
    Configuration menu
    Copy the full SHA
    8843937 View commit details
    Browse the repository at this point in the history

  31. Compute max cost matrix

    @michalk8
    michalk8 committed Oct 10, 2024
    Configuration menu
    Copy the full SHA
    83c9960 View commit details
    Browse the repository at this point in the history

Commits on Oct 11, 2024

  1. Update error message

    @michalk8
    michalk8 committed Oct 11, 2024
    Configuration menu
    Copy the full SHA
    69a9599 View commit details
    Browse the repository at this point in the history

  2. Remove TODO

    @michalk8
    michalk8 committed Oct 11, 2024
    Configuration menu
    Copy the full SHA
    6667abd View commit details
    Browse the repository at this point in the history

  3. Flatten out axes

    @michalk8
    michalk8 committed Oct 11, 2024
    Configuration menu
    Copy the full SHA
    ac9b928 View commit details
    Browse the repository at this point in the history

  4. Fix missing cross terms in the costs

    @michalk8
    michalk8 committed Oct 11, 2024
    Configuration menu
    Copy the full SHA
    c113946 View commit details
    Browse the repository at this point in the history

  5. Fix geom tests

    @michalk8
    michalk8 committed Oct 11, 2024
    Configuration menu
    Copy the full SHA
    75d9e7a View commit details
    Browse the repository at this point in the history

  6. Fix dtype

    @michalk8
    michalk8 committed Oct 11, 2024
    Configuration menu
    Copy the full SHA
    44eb5a8 View commit details
    Browse the repository at this point in the history

  7. Start implementing transport functions

    @michalk8
    michalk8 committed Oct 11, 2024
    Configuration menu
    Copy the full SHA
    cbb4ea0 View commit details
    Browse the repository at this point in the history

  8. Implement online transport functions

    @michalk8
    michalk8 committed Oct 11, 2024
    Configuration menu
    Copy the full SHA
    e8bb1b5 View commit details
    Browse the repository at this point in the history

  9. Fix solver tests

    @michalk8
    michalk8 committed Oct 11, 2024
    Configuration menu
    Copy the full SHA
    7d4001e View commit details
    Browse the repository at this point in the history

  10. Fix Bures test

    @michalk8
    michalk8 committed Oct 11, 2024
    Configuration menu
    Copy the full SHA
    8941224 View commit details
    Browse the repository at this point in the history

  11. Don't use pairwise in tests

    @michalk8
    michalk8 committed Oct 11, 2024
    Configuration menu
    Copy the full SHA
    a565e09 View commit details
    Browse the repository at this point in the history

  12. Update notebook that uses norm

    @michalk8
    michalk8 committed Oct 11, 2024
    Configuration menu
    Copy the full SHA
    1533324 View commit details
    Browse the repository at this point in the history

  13. Fix bug in UnbalancedBures

    @michalk8
    michalk8 committed Oct 11, 2024
    Configuration menu
    Copy the full SHA
    3e7ff8b View commit details
    Browse the repository at this point in the history

  14. Rename pairwise -> __call__

    @michalk8
    michalk8 committed Oct 11, 2024
    Configuration menu
    Copy the full SHA
    b815fbc View commit details
    Browse the repository at this point in the history

  15. Remove old shape code

    @michalk8
    michalk8 committed Oct 11, 2024
    Configuration menu
    Copy the full SHA
    739afde View commit details
    Browse the repository at this point in the history

  16. Always instantiate the cost for online

    @michalk8
    michalk8 committed Oct 11, 2024
    Configuration menu
    Copy the full SHA
    0d7f6ae View commit details
    Browse the repository at this point in the history

  17. Remove old TODO

    @michalk8
    michalk8 committed Oct 11, 2024
    Configuration menu
    Copy the full SHA
    4863fcf View commit details
    Browse the repository at this point in the history

  18. Extract _apply_cost_to_vec_fast

    @michalk8
    michalk8 committed Oct 11, 2024
    Configuration menu
    Copy the full SHA
    4aa4c6b View commit details
    Browse the repository at this point in the history

  19. Update max cost in LRCGeom

    @michalk8
    michalk8 committed Oct 11, 2024
    Configuration menu
    Copy the full SHA
    8511073 View commit details
    Browse the repository at this point in the history

  20. Fix test, use more multi_dot

    @michalk8
    michalk8 committed Oct 11, 2024
    Configuration menu
    Copy the full SHA
    47462d2 View commit details
    Browse the repository at this point in the history

  21. Remove batch_size from LRCGeometry

    @michalk8
    michalk8 committed Oct 11, 2024
    Configuration menu
    Copy the full SHA
    05630a8 View commit details
    Browse the repository at this point in the history

Commits on Oct 15, 2024

  1. Add better warning error

    @michalk8
    michalk8 committed Oct 15, 2024
    Configuration menu
    Copy the full SHA
    0994d7a View commit details
    Browse the repository at this point in the history

  2. Reorder properties

    @michalk8
    michalk8 committed Oct 15, 2024
    Configuration menu
    Copy the full SHA
    5d88ad4 View commit details
    Browse the repository at this point in the history

  3. Add docs to batched_vmap

    @michalk8
    michalk8 committed Oct 15, 2024
    Configuration menu
    Copy the full SHA
    f8143fc View commit details
    Browse the repository at this point in the history

  4. Start adding tests

    @michalk8
    michalk8 committed Oct 15, 2024
    Configuration menu
    Copy the full SHA
    a82688c View commit details
    Browse the repository at this point in the history

  5. Reorder functions in test

    @michalk8
    michalk8 committed Oct 15, 2024
    Configuration menu
    Copy the full SHA
    1d2d12d View commit details
    Browse the repository at this point in the history

  6. Fix axes, add a test

    @michalk8
    michalk8 committed Oct 15, 2024
    Configuration menu
    Copy the full SHA
    44b1126 View commit details
    Browse the repository at this point in the history

  7. Update test fn

    @michalk8
    michalk8 committed Oct 15, 2024
    Configuration menu
    Copy the full SHA
    889f81f View commit details
    Browse the repository at this point in the history

  8. Move out assert

    @michalk8
    michalk8 committed Oct 15, 2024
    Configuration menu
    Copy the full SHA
    b16d5a8 View commit details
    Browse the repository at this point in the history

  9. Dont canon out_axes

    @michalk8
    michalk8 committed Oct 15, 2024
    Configuration menu
    Copy the full SHA
    c984a43 View commit details
    Browse the repository at this point in the history

  10. Check max traces

    @michalk8
    michalk8 committed Oct 15, 2024
    Configuration menu
    Copy the full SHA
    4426994 View commit details
    Browse the repository at this point in the history

  11. Test memory of batched vmap

    @michalk8
    michalk8 committed Oct 15, 2024
    Configuration menu
    Copy the full SHA
    5e5125b View commit details
    Browse the repository at this point in the history

  12. Install typing_extensions

    @michalk8
    michalk8 committed Oct 15, 2024
    Configuration menu
    Copy the full SHA
    cb31db7 View commit details
    Browse the repository at this point in the history

  13. Merge branch 'main' into feature/batched-vmap

    @michalk8
    michalk8 committed Oct 15, 2024
    Configuration menu
    Copy the full SHA
    57bf9ca View commit details
    Browse the repository at this point in the history

  14. Remove . from description

    @michalk8
    michalk8 committed Oct 15, 2024
    Configuration menu
    Copy the full SHA
    721eca9 View commit details
    Browse the repository at this point in the history

  15. Add more out_axes tests

    @michalk8
    michalk8 committed Oct 15, 2024
    Configuration menu
    Copy the full SHA
    f9a41bd View commit details
    Browse the repository at this point in the history

  16. Add in_axes test

    @michalk8
    michalk8 committed Oct 15, 2024
    Configuration menu
    Copy the full SHA
    78003d9 View commit details
    Browse the repository at this point in the history

  17. Fix negative axes

    @michalk8
    michalk8 committed Oct 15, 2024
    Configuration menu
    Copy the full SHA
    9e1ae03 View commit details
    Browse the repository at this point in the history

Commits on Oct 16, 2024

  1. Increase memory limit in the test

    @michalk8
    michalk8 committed Oct 16, 2024
    Configuration menu
    Copy the full SHA
    427b5ec View commit details
    Browse the repository at this point in the history

  2. Add in_axes pytree test

    @michalk8
    michalk8 committed Oct 16, 2024
    Configuration menu
    Copy the full SHA
    fff0ce6 View commit details
    Browse the repository at this point in the history

  3. Remove old warnings filters

    @michalk8
    michalk8 committed Oct 16, 2024
    Configuration menu
    Copy the full SHA
    f72abf1 View commit details
    Browse the repository at this point in the history

  4. Update fixtures

    @michalk8
    michalk8 committed Oct 16, 2024
    Configuration menu
    Copy the full SHA
    b19ff4b View commit details
    Browse the repository at this point in the history

  5. Update SqEucl cost.

    @michalk8
    michalk8 committed Oct 16, 2024
    Configuration menu
    Copy the full SHA
    462c630 View commit details
    Browse the repository at this point in the history

  6. Update docstrings

    @michalk8
    michalk8 committed Oct 16, 2024
    Configuration menu
    Copy the full SHA
    babb095 View commit details
    Browse the repository at this point in the history

  7. Remove unused imports from the docs

    @michalk8
    michalk8 committed Oct 16, 2024
    Configuration menu
    Copy the full SHA
    87df731 View commit details
    Browse the repository at this point in the history

  8. Revert test pre-commits

    @michalk8
    michalk8 committed Oct 16, 2024
    Configuration menu
    Copy the full SHA
    07dff82 View commit details
    Browse the repository at this point in the history

  9. Fix ICNN init notebook 

    Was broken by #551
    @michalk8
    michalk8 committed Oct 16, 2024
    Configuration menu
    Copy the full SHA
    5450808 View commit details
    Browse the repository at this point in the history

  10. Improve error message

    @michalk8
    michalk8 committed Oct 16, 2024
    Configuration menu
    Copy the full SHA
    e390e64 View commit details
    Browse the repository at this point in the history