Use convert_element_type instead of device_put_raw. by pschuh · Pull Request #6400 · jax-ml/jax

pschuh · 2021-04-09T22:53:58Z

No description provided.

jakevdp · 2021-04-10T00:01:40Z

Hi - out of curiosity, what was the impetus for this change? I ask because this kind of change tends to have subtle unintented consequences, so I'd prefer to avoid it unless there is good reason.

mattjj · 2021-04-10T01:44:15Z

Also, this code was written this way for performance, so we'd want to check some benchmarks for any change! #3350 has some basic microbenchmarks for jnp.array. (Now that we have some basic benchmark infrastructure, we should add this kind of stuff! )

jax/_src/numpy/lax_numpy.py

mattjj

Great find, Parker! (See jax-dev chat thread for more discussion.)

mattjj · 2021-04-10T04:34:27Z

I ran the simple timeit jnp.array([0] * int(1e6)) benchmark from #3350 and there was no change (214 ms ± 1.95 on master, 213 ms ± 1.88 on this branch).

For reference: jax-ml/jax#6400

For reference: jax-ml/jax#6400.

For reference: jax-ml/jax#6400. PiperOrigin-RevId: 369466397

For reference: jax-ml/jax#6400. PiperOrigin-RevId: 369471455

also add a test for not performing H2D transfers while tracing jnp.array

mattjj · 2021-04-22T19:47:52Z

@jakevdp the main motivation is to avoid host-to-device (H2D) transfers during jit tracing, very much analogous to #6014. See the jax-dev chat thread for more discussion.

With jax-ml/jax#6400, jnp.array delays resolving the axis name until after tracing is complete. PiperOrigin-RevId: 370356687

With jax-ml/jax#6400, jnp.array delays resolving the axis name until after tracing is complete. PiperOrigin-RevId: 370798828

For reference: jax-ml/jax#6400. PiperOrigin-RevId: 372376815

google-cla · 2021-05-06T20:18:11Z

All (the pull request submitter and all commit authors) CLAs are signed, but one or more commits were authored or co-authored by someone other than the pull request submitter.

We need to confirm that all authors are ok with their commits being contributed to this project. Please have them confirm that by leaving a comment that contains only @googlebot I consent. in this pull request.

Note to project maintainer: There may be cases where the author cannot leave a comment, or the comment is not properly detected as consent. In those cases, you can manually confirm consent of the commit author(s), and set the cla label to yes (if enabled on your project).

ℹ️ Googlers: Go here for more info.

mattjj · 2021-05-06T20:24:15Z

@googlebot I consent.

google-cla · 2021-05-06T20:24:23Z

All (the pull request submitter and all commit authors) CLAs are signed, but one or more commits were authored or co-authored by someone other than the pull request submitter.

We need to confirm that all authors are ok with their commits being contributed to this project. Please have them confirm that by leaving a comment that contains only @googlebot I consent. in this pull request.

Note to project maintainer: There may be cases where the author cannot leave a comment, or the comment is not properly detected as consent. In those cases, you can manually confirm consent of the commit author(s), and set the cla label to yes (if enabled on your project).

ℹ️ Googlers: Go here for more info.

This PR changes `jax.numpy.array()` to avoid creating any on-device arrays during tracing. As a consequence, calls to `jnp.array()` in a traced context, such as `jax.jit` will always be staged into the trace. This change may break code that depends on the current (undocumented and unintentional) behavior of `jnp.array()` to perform shape or index calculations that must be known statically (at trace time). The workaround for such cases is to use classic NumPy to perform shape/index calculations. PiperOrigin-RevId: 395998020

This PR changes `jax.numpy.array()` to avoid creating any on-device arrays during tracing. As a consequence, calls to `jnp.array()` in a traced context, such as `jax.jit` will always be staged into the trace. This change may break code that depends on the current (undocumented and unintentional) behavior of `jnp.array()` to perform shape or index calculations that must be known statically (at trace time). The workaround for such cases is to use classic NumPy to perform shape/index calculations. PiperOrigin-RevId: 398008511

For reference: jax-ml/jax#6400. PiperOrigin-RevId: 369471455

With jax-ml/jax#6400, jnp.array delays resolving the axis name until after tracing is complete. PiperOrigin-RevId: 370798828

google-cla bot added the cla: yes label Apr 9, 2021

mattjj self-assigned this Apr 10, 2021

mattjj self-requested a review April 10, 2021 03:37

pschuh commented Apr 10, 2021

View reviewed changes

jax/_src/numpy/lax_numpy.py Outdated Show resolved Hide resolved

mattjj approved these changes Apr 10, 2021

View reviewed changes

mattjj force-pushed the convert_element branch from e981c29 to 4cc07d7 Compare April 10, 2021 04:32

mattjj added the pull ready Ready for copybara import and testing label Apr 10, 2021

pschuh added a commit to pschuh/flax that referenced this pull request Apr 20, 2021

Fix places that were passing jnp.array where a shape was expected.

05bd185

For reference: jax-ml/jax#6400

pschuh mentioned this pull request Apr 20, 2021

Fix places that were passing jnp.array where a shape was expected. google/flax#1258

Merged

pschuh added a commit to pschuh/trax that referenced this pull request Apr 20, 2021

Fix places that were passing jnp.array where a shape was expected.

a6ac5aa

For reference: jax-ml/jax#6400.

pschuh mentioned this pull request Apr 20, 2021

Fix places that were passing jnp.array where a shape was expected. google/trax#1595

Closed

copybara-service bot pushed a commit to google/trax that referenced this pull request Apr 20, 2021

Fix places that were passing jnp.array where a shape was expected.

3a9859d

For reference: jax-ml/jax#6400. PiperOrigin-RevId: 369466397

copybara-service bot mentioned this pull request Apr 20, 2021

Fix places that were passing jnp.array where a shape was expected. google/trax#1597

Merged

copybara-service bot pushed a commit to google/trax that referenced this pull request Apr 20, 2021

Fix places that were passing jnp.array where a shape was expected.

1792bf5

For reference: jax-ml/jax#6400. PiperOrigin-RevId: 369466397

copybara-service bot pushed a commit to google/trax that referenced this pull request Apr 20, 2021

Fix places that were passing jnp.array where a shape was expected.

28743ee

For reference: jax-ml/jax#6400. PiperOrigin-RevId: 369471455

mattjj added pull ready Ready for copybara import and testing and removed pull ready Ready for copybara import and testing labels Apr 22, 2021

pschuh and others added 2 commits April 22, 2021 12:46

Use convert_element_type insteadof device_put_raw.

99b755f

prune trivial convert_element_types from jaxprs

ba9233b

also add a test for not performing H2D transfers while tracing jnp.array

mattjj force-pushed the convert_element branch from 4cc07d7 to ba9233b Compare April 22, 2021 19:46

copybara-service bot pushed a commit to google/trax that referenced this pull request Apr 25, 2021

Change jnp.array to np.array.

d8f9c26

With jax-ml/jax#6400, jnp.array delays resolving the axis name until after tracing is complete. PiperOrigin-RevId: 370356687

copybara-service bot mentioned this pull request Apr 25, 2021

Change jnp.array(1.0) to 1. google/trax#1604

Merged

jnp.asarray -> _device_put_raw for call_tf.py

2c35b1a

copybara-service bot pushed a commit to google/trax that referenced this pull request Apr 25, 2021

Change jnp.array(1.0) to 1.

3d087ed

With jax-ml/jax#6400, jnp.array delays resolving the axis name until after tracing is complete. PiperOrigin-RevId: 370356687

copybara-service bot pushed a commit to google/trax that referenced this pull request Apr 28, 2021

Change jnp.array(1.0) to 1.

5d1baa1

With jax-ml/jax#6400, jnp.array delays resolving the axis name until after tracing is complete. PiperOrigin-RevId: 370356687

copybara-service bot pushed a commit to google/trax that referenced this pull request Apr 28, 2021

Change jnp.array(1.0) to 1.

2e484f5

With jax-ml/jax#6400, jnp.array delays resolving the axis name until after tracing is complete. PiperOrigin-RevId: 370798828

tensorflow-copybara pushed a commit to google-research/google-research that referenced this pull request May 6, 2021

Fix places that were passing jnp.array where a shape was expected.

3fac129

For reference: jax-ml/jax#6400. PiperOrigin-RevId: 372376815

mattjj added pull ready Ready for copybara import and testing and removed pull ready Ready for copybara import and testing labels May 6, 2021

Merge branch 'master' into convert_element

9d3e535

google-cla bot added cla: no and removed cla: yes labels May 6, 2021

mattjj added cla: yes and removed cla: no labels May 6, 2021

google-cla bot added cla: no and removed cla: yes labels May 6, 2021

mattjj added cla: yes and removed cla: no labels May 6, 2021

copybara-service bot merged commit d0aa875 into jax-ml:master May 6, 2021

pschuh mentioned this pull request May 7, 2021

Revert "Use convert_element_type instead of device_put_raw." #6675

Merged

copybara-service bot mentioned this pull request Sep 21, 2021

Attempt to land https://github.com/google/jax/pull/6400 again. #7967

Merged

mmarcinmichal pushed a commit to mmarcinmichal/trax that referenced this pull request Jan 13, 2026

Fix places that were passing jnp.array where a shape was expected.

81564be

For reference: jax-ml/jax#6400. PiperOrigin-RevId: 369471455

mmarcinmichal pushed a commit to mmarcinmichal/trax that referenced this pull request Jan 13, 2026

Change jnp.array(1.0) to 1.

818984a

With jax-ml/jax#6400, jnp.array delays resolving the axis name until after tracing is complete. PiperOrigin-RevId: 370798828

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Use convert_element_type instead of device_put_raw.#6400

Use convert_element_type instead of device_put_raw.#6400
copybara-service[bot] merged 4 commits intojax-ml:masterfrom
pschuh:convert_element

pschuh commented Apr 9, 2021

Uh oh!

jakevdp commented Apr 10, 2021

Uh oh!

mattjj commented Apr 10, 2021

Uh oh!

Uh oh!

mattjj left a comment

Uh oh!

mattjj commented Apr 10, 2021

Uh oh!

mattjj commented Apr 22, 2021

Uh oh!

google-cla bot commented May 6, 2021

Uh oh!

mattjj commented May 6, 2021 •

edited

Loading

Uh oh!

google-cla bot commented May 6, 2021

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

pschuh commented Apr 9, 2021

Uh oh!

jakevdp commented Apr 10, 2021

Uh oh!

mattjj commented Apr 10, 2021

Uh oh!

Uh oh!

mattjj left a comment

Choose a reason for hiding this comment

Uh oh!

mattjj commented Apr 10, 2021

Uh oh!

mattjj commented Apr 22, 2021

Uh oh!

google-cla bot commented May 6, 2021

Uh oh!

mattjj commented May 6, 2021 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

google-cla bot commented May 6, 2021

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

mattjj commented May 6, 2021 •

edited

Loading