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

Implement lax.map #1118

Merged
merged 2 commits into from
Aug 5, 2019
Merged

Implement lax.map #1118

merged 2 commits into from
Aug 5, 2019

Conversation

shoyer
Copy link
Collaborator

@shoyer shoyer commented Aug 5, 2019

Fixes GH-1113

@shoyer shoyer mentioned this pull request Aug 5, 2019
Closed
ibab
ibab reviewed Aug 5, 2019
View reviewed changes
jax/lax/lax_control_flow.py Outdated
def map(f, xs):
return np.stack([f(x) for x in xs])

Like ``scan``, ``map`` is implemented in terms of JAX primtivies so many of
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

primtivies -> primitives

mattjj
mattjj approved these changes Aug 5, 2019
View reviewed changes
Copy link
Collaborator

@mattjj mattjj left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM! Will merge after tests pass.

(We could even imagine having a version of api.vmap that uses a combination of vectorization and this kind of "old-fashioned" looping to fit vmapped computations into a memory budget...)

@mattjj mattjj merged commit f62c2e4 into jax-ml:master Aug 5, 2019
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@ibab ibab ibab left review comments

@mattjj mattjj mattjj approved these changes

Assignees
No one assigned
Labels
cla: yes
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.

lax.map
4 participants
@shoyer @ibab @mattjj @googlebot