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

How well does tree-math support computation on multiple devices? #10

Closed
connection-on-fiber-bundles opened this issue Apr 21, 2022 · 1 comment
Closed

How well does tree-math support computation on multiple devices? #10

connection-on-fiber-bundles opened this issue Apr 21, 2022 · 1 comment

Comments

@connection-on-fiber-bundles
Copy link

Wondering how well tree-math supports computation on multiple devices?

Let's say we have a pytree of tensors of different dimensions and want to perform some operations on each of them with tree-math, can we distribute those tasks to multiple devices (GPU, for instance)?
@shoyer
Copy link
Member

shoyer commented Apr 25, 2022 •
edited
Loading

This should not be a problem. Tree-Math is entirely agnostic to JAX's multi-device APIs. It's just syntactic sugar for jax.tree_util.tree_map.

@shoyer shoyer closed this as completed May 18, 2023
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Assignees
No one assigned
Labels
None yet
Projects
None yet
Milestone
No milestone
Development

No branches or pull requests
2 participants
@shoyer @connection-on-fiber-bundles