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

Fix axis_index inside nested pmaps #4378

Merged
merged 2 commits into from
Sep 22, 2020
Merged

Fix axis_index inside nested pmaps #4378

merged 2 commits into from
Sep 22, 2020

Commits on Sep 22, 2020

  1. Fix axis_index inside nested pmaps 

    The previous translation rule has assumed that `axis_index` is always
taken over the outermost axis in the `axis_env`, and was always producing
the same output, no matter which axis has been specified. This fixes the
translation rule to start taking the `axis_name` into account.

Additionally, this adds support for querying the index along multiple
axes, which will be useful for `gmap`.
    @apaszke
    apaszke committed Sep 22, 2020
    Configuration menu
    Copy the full SHA
    332a9ba View commit details
    Browse the repository at this point in the history

  2. Fix a faulty soft_pmap rule for axis_index 

    The rule didn't specify the precision for the `np.arange` constant,
which caused an accidental dtype promotion in X64 mode. Previously the
error has luckicly been hidden behind a coerction that followed
`axis_index` in that test, but the new implementation has surfaced it.
    @apaszke
    apaszke committed Sep 22, 2020
    Configuration menu
    Copy the full SHA
    8ac19c7 View commit details
    Browse the repository at this point in the history