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.

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

Replace indexed_shape by a version that uses jit #519

Closed
wants to merge 3 commits into from

Conversation

Michael-T-McCann
Copy link
Contributor

@Michael-T-McCann Michael-T-McCann commented May 15, 2024

Followup to #517 that is both fewer lines of code and more flexible in that any indexing expression that is valid for jax arrays should automatically work.

@bwohlberg
Copy link
Collaborator

Clever use of jax.numpy.empty. What's the computational cost of redoing the jit on each call? Have you checked that it's not considerably slower than the version it replaces?

@Michael-T-McCann
Copy link
Contributor Author

Timing with

import scico
import numpy as np
%timeit -n 1 -r 1 scico.numpy.util.indexed_shape((1000, 1000, 1000, 1000), (np.newaxis, ..., slice(0, 10), slice(15, 0, -1)))

Proposed version: ~10 ms, existing version, ~10 µs. So we we 1000x slower but these are small numbers.

@bwohlberg
Copy link
Collaborator

I hesitate to merge this PR due to the difference in execution time, and also because the new version of indexed_shape will consume potentially large amounts of memory if JIT is disabled for debugging. I propose the compromise in branch brendt/index_shape: keep the original function, and add this one as a potentially useful reference version that is not used by other package components.

@bwohlberg bwohlberg added enhancement New feature or request improvement Improvement of existing code, including addressing of omissions or inconsistencies and removed enhancement New feature or request labels May 22, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
improvement Improvement of existing code, including addressing of omissions or inconsistencies
Projects
None yet
Development

Successfully merging this pull request may close these issues.

None yet

2 participants