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

JAX internal functions must take a non-empty core.DebugInfo. #26480

Open
gnecula opened this issue Feb 12, 2025 · 0 comments
Open

JAX internal functions must take a non-empty core.DebugInfo. #26480

gnecula opened this issue Feb 12, 2025 · 0 comments
Assignees
@gnecula
Labels
better_errors Improve the error reporting

Comments

@gnecula
Copy link
Collaborator

gnecula commented Feb 12, 2025 •
edited
Loading

Description

As part of improving JAX error messages we ensure that a core.DebugInfo object is passed to all calls to linear_util.wrap_init and to the core.Jaxpr constructor. We have made the necessary changes into the JAX core code base (e.g., #26078, #26313, #26348, and several other PRs).

There are still a number of other projects that invoke directly these functions (and other helper functions that now must take core.DebugInfo). Note that only JAX-internal functions take a core.DebugInfo; this change does not affect projects that use only public APIs.

For a limited time, a DeprecationWarning is printed if jax.extend.linear_util.wrap_init is used without debugging info.

To obtain a core.DebugInfo, one can call api_util.debug_info, or retrieve the debug_info fields of lu.WrappedFun or core.Jaxpr.

System info (python version, jaxlib version, accelerator, etc.)

N/A
@gnecula gnecula added the bug Something isn't working label Feb 12, 2025
@gnecula gnecula self-assigned this Feb 12, 2025
@gnecula gnecula added better_errors Improve the error reporting and removed bug Something isn't working labels Feb 12, 2025
copybara-service bot pushed a commit to jax-ml/oryx that referenced this issue Feb 12, 2025
@gnecula
[better_errors] Make it explicit that debug_info is not None.
e64df5c 
Now all internal uses of lu.wrap_init and core.Jaxpr are with actual
debug info. This enables us to clean up the type declarations and
to remove the checks whether debug_info is present.

For usage outside of the JAX internals, we change
`jax.extend.linear_util.wrap_init` to be usable without debug_info,
for temporary backwards compatibility. We emit a deprecation
warning and fill-in some fake debugging info.

See jax-ml/jax#26480 for more details.

PiperOrigin-RevId: 725512692
@copybara-service copybara-service bot mentioned this issue Feb 12, 2025
Open
copybara-service bot pushed a commit that referenced this issue Feb 12, 2025
@gnecula @Google-ML-Automation
[better_errors] Make it explicit that debug_info is not None.
b5a745d 
Now all internal uses of lu.wrap_init and core.Jaxpr are with actual
debug info. This enables us to clean up the type declarations and
to remove the checks whether debug_info is present.

For usage outside of the JAX internals, we change
`jax.extend.linear_util.wrap_init` to be usable without debug_info,
for temporary backwards compatibility. We emit a deprecation
warning and fill-in some fake debugging info.

See #26480 for more details.

PiperOrigin-RevId: 725512692
@copybara-service copybara-service bot mentioned this issue Feb 12, 2025
Closed
copybara-service bot pushed a commit to jax-ml/oryx that referenced this issue Feb 12, 2025
@gnecula
Update oryx for JAX debug_info
2463d10 
See jax-ml/jax#26480.

PiperOrigin-RevId: 725953448
@copybara-service copybara-service bot mentioned this issue Feb 12, 2025
Merged
copybara-service bot pushed a commit that referenced this issue Feb 12, 2025
@gnecula @Google-ML-Automation
[better_errors] Make it explicit that debug_info is not None.
95a87b2 
Now all internal uses of lu.wrap_init and core.Jaxpr are with actual
debug info. This enables us to clean up the type declarations and
to remove the checks whether debug_info is present.

For usage outside of the JAX internals, we change
`jax.extend.linear_util.wrap_init` to be usable without debug_info,
for temporary backwards compatibility. We emit a deprecation
warning and fill-in some fake debugging info.

See #26480 for more details.

PiperOrigin-RevId: 725512692
copybara-service bot pushed a commit to jax-ml/oryx that referenced this issue Feb 12, 2025
@gnecula
[better_errors] Make it explicit that debug_info is not None.
3856e38 
Now all internal uses of lu.wrap_init and core.Jaxpr are with actual
debug info. This enables us to clean up the type declarations and
to remove the checks whether debug_info is present.

For usage outside of the JAX internals, we change
`jax.extend.linear_util.wrap_init` to be usable without debug_info,
for temporary backwards compatibility. We emit a deprecation
warning and fill-in some fake debugging info.

See jax-ml/jax#26480 for more details.

PiperOrigin-RevId: 725512692
copybara-service bot pushed a commit that referenced this issue Feb 12, 2025
@gnecula @Google-ML-Automation
[better_errors] Make it explicit that debug_info is not None.
6997e5c 
Now all internal uses of lu.wrap_init and core.Jaxpr are with actual
debug info. This enables us to clean up the type declarations and
to remove the checks whether debug_info is present.

For usage outside of the JAX internals, we change
`jax.extend.linear_util.wrap_init` to be usable without debug_info,
for temporary backwards compatibility. We emit a deprecation
warning and fill-in some fake debugging info.

See #26480 for more details.

PiperOrigin-RevId: 725512692
copybara-service bot pushed a commit to jax-ml/oryx that referenced this issue Feb 12, 2025
@gnecula
Update oryx for JAX debug_info
1053cf7 
See jax-ml/jax#26480.

PiperOrigin-RevId: 725953448
copybara-service bot pushed a commit that referenced this issue Feb 12, 2025
@gnecula @Google-ML-Automation
Update oryx for JAX debug_info
02ad3d9 
See #26480.

PiperOrigin-RevId: 725953448
@copybara-service copybara-service bot mentioned this issue Feb 12, 2025
Merged
copybara-service bot pushed a commit to jax-ml/oryx that referenced this issue Feb 12, 2025
@gnecula
Update oryx for JAX debug_info
f4a02e3 
See jax-ml/jax#26480.

PiperOrigin-RevId: 725953448
copybara-service bot pushed a commit that referenced this issue Feb 12, 2025
@gnecula @Google-ML-Automation
Update oryx for JAX debug_info
76c7df4 
See #26480.

PiperOrigin-RevId: 725953448
copybara-service bot pushed a commit to jax-ml/oryx that referenced this issue Feb 12, 2025
@gnecula
[better_errors] Make it explicit that debug_info is not None.
8a54cef 
Now all internal uses of lu.wrap_init and core.Jaxpr are with actual
debug info. This enables us to clean up the type declarations and
to remove the checks whether debug_info is present.

For usage outside of the JAX internals, we change
`jax.extend.linear_util.wrap_init` to be usable without debug_info,
for temporary backwards compatibility. We emit a deprecation
warning and fill-in some fake debugging info.

See jax-ml/jax#26480 for more details.

PiperOrigin-RevId: 725512692
copybara-service bot pushed a commit that referenced this issue Feb 12, 2025
@gnecula @Google-ML-Automation
[better_errors] Make it explicit that debug_info is not None.
9d32061 
Now all internal uses of lu.wrap_init and core.Jaxpr are with actual
debug info. This enables us to clean up the type declarations and
to remove the checks whether debug_info is present.

For usage outside of the JAX internals, we change
`jax.extend.linear_util.wrap_init` to be usable without debug_info,
for temporary backwards compatibility. We emit a deprecation
warning and fill-in some fake debugging info.

See #26480 for more details.

PiperOrigin-RevId: 725512692
copybara-service bot pushed a commit to jax-ml/oryx that referenced this issue Feb 12, 2025
@gnecula
[better_errors] Make it explicit that debug_info is not None.
ace2481 
Now all internal uses of lu.wrap_init and core.Jaxpr are with actual
debug info. This enables us to clean up the type declarations and
to remove the checks whether debug_info is present.

For usage outside of the JAX internals, we change
`jax.extend.linear_util.wrap_init` to be usable without debug_info,
for temporary backwards compatibility. We emit a deprecation
warning and fill-in some fake debugging info.

See jax-ml/jax#26480 for more details.

PiperOrigin-RevId: 725512692
copybara-service bot pushed a commit that referenced this issue Feb 12, 2025
@gnecula @Google-ML-Automation
[better_errors] Make it explicit that debug_info is not None.
c64cfb6 
Now all internal uses of lu.wrap_init and core.Jaxpr are with actual
debug info. This enables us to clean up the type declarations and
to remove the checks whether debug_info is present.

For usage outside of the JAX internals, we change
`jax.extend.linear_util.wrap_init` to be usable without debug_info,
for temporary backwards compatibility. We emit a deprecation
warning and fill-in some fake debugging info.

See #26480 for more details.

PiperOrigin-RevId: 725512692
copybara-service bot pushed a commit to jax-ml/oryx that referenced this issue Feb 12, 2025
@gnecula
Update oryx for JAX debug_info
2a24555 
See jax-ml/jax#26480.

PiperOrigin-RevId: 725953448
copybara-service bot pushed a commit that referenced this issue Feb 12, 2025
@gnecula @Google-ML-Automation
Update oryx for JAX debug_info
5dcb1a3 
See #26480.

PiperOrigin-RevId: 725953448
copybara-service bot pushed a commit to jax-ml/oryx that referenced this issue Feb 12, 2025
@gnecula
[better_errors] Make it explicit that debug_info is not None.
96007cc 
Now all internal uses of lu.wrap_init and core.Jaxpr are with actual
debug info. This enables us to clean up the type declarations and
to remove the checks whether debug_info is present.

For usage outside of the JAX internals, we change
`jax.extend.linear_util.wrap_init` to be usable without debug_info,
for temporary backwards compatibility. We emit a deprecation
warning and fill-in some fake debugging info.

See jax-ml/jax#26480 for more details.

PiperOrigin-RevId: 725512692
copybara-service bot pushed a commit that referenced this issue Feb 12, 2025
@gnecula @Google-ML-Automation
[better_errors] Make it explicit that debug_info is not None.
9690625 
Now all internal uses of lu.wrap_init and core.Jaxpr are with actual
debug info. This enables us to clean up the type declarations and
to remove the checks whether debug_info is present.

For usage outside of the JAX internals, we change
`jax.extend.linear_util.wrap_init` to be usable without debug_info,
for temporary backwards compatibility. We emit a deprecation
warning and fill-in some fake debugging info.

See #26480 for more details.

PiperOrigin-RevId: 725512692
copybara-service bot pushed a commit to jax-ml/oryx that referenced this issue Feb 12, 2025
@gnecula
Update oryx for JAX debug_info
41577a1 
See jax-ml/jax#26480.

PiperOrigin-RevId: 725953448
copybara-service bot pushed a commit that referenced this issue Feb 12, 2025
@gnecula @Google-ML-Automation
Update oryx for JAX debug_info
219031a 
See #26480.

PiperOrigin-RevId: 725953448
copybara-service bot pushed a commit to jax-ml/oryx that referenced this issue Feb 12, 2025
@gnecula
Update oryx for JAX debug_info
edc35b3 
See jax-ml/jax#26480.

PiperOrigin-RevId: 725953448
copybara-service bot pushed a commit that referenced this issue Feb 12, 2025
@gnecula @Google-ML-Automation
Update oryx for JAX debug_info
9259efb 
See #26480.

PiperOrigin-RevId: 725953448
copybara-service bot pushed a commit to jax-ml/oryx that referenced this issue Feb 13, 2025
@gnecula
Update oryx for JAX debug_info
d75d6b3 
See jax-ml/jax#26480.

PiperOrigin-RevId: 725953448
copybara-service bot pushed a commit that referenced this issue Feb 13, 2025
@gnecula @Google-ML-Automation
Update oryx for JAX debug_info
c85876a 
See #26480.

PiperOrigin-RevId: 725953448
copybara-service bot pushed a commit to jax-ml/oryx that referenced this issue Feb 13, 2025
@gnecula
Update oryx for JAX debug_info
73d2f98 
See jax-ml/jax#26480.

PiperOrigin-RevId: 725953448
copybara-service bot pushed a commit that referenced this issue Feb 13, 2025
@gnecula @Google-ML-Automation
Update oryx for JAX debug_info
931029f 
See #26480.

PiperOrigin-RevId: 725953448
copybara-service bot pushed a commit to jax-ml/oryx that referenced this issue Feb 13, 2025
@gnecula
Update oryx for JAX debug_info
342d633 
See jax-ml/jax#26480.

PiperOrigin-RevId: 725953448
copybara-service bot pushed a commit that referenced this issue Feb 13, 2025
@gnecula @Google-ML-Automation
Update oryx for JAX debug_info
b542f1b 
See #26480.

PiperOrigin-RevId: 725953448
copybara-service bot pushed a commit to jax-ml/oryx that referenced this issue Feb 13, 2025
@gnecula
Update oryx for JAX debug_info
a21730c 
See jax-ml/jax#26480.

PiperOrigin-RevId: 725953448
copybara-service bot pushed a commit that referenced this issue Feb 13, 2025
@gnecula @Google-ML-Automation
Update oryx for JAX debug_info
53e08d8 
See #26480.

PiperOrigin-RevId: 725953448
copybara-service bot pushed a commit to jax-ml/oryx that referenced this issue Feb 13, 2025
@gnecula
Update oryx for JAX debug_info
3f15505 
See jax-ml/jax#26480.

PiperOrigin-RevId: 725953448
copybara-service bot pushed a commit that referenced this issue Feb 13, 2025
@gnecula @Google-ML-Automation
Update oryx for JAX debug_info
e062c6b 
See #26480.

PiperOrigin-RevId: 725953448
copybara-service bot pushed a commit to jax-ml/oryx that referenced this issue Feb 13, 2025
@gnecula
Update oryx for JAX debug_info
edaf49a 
See jax-ml/jax#26480.

PiperOrigin-RevId: 725953448
copybara-service bot pushed a commit that referenced this issue Feb 13, 2025
@gnecula @Google-ML-Automation
Update oryx for JAX debug_info
19ea81b 
See #26480.

PiperOrigin-RevId: 725953448
copybara-service bot pushed a commit that referenced this issue Feb 13, 2025
@gnecula @Google-ML-Automation
Update oryx for JAX debug_info
7161cad 
See #26480.

PiperOrigin-RevId: 726437113
copybara-service bot pushed a commit to jax-ml/oryx that referenced this issue Feb 13, 2025
@gnecula
Update oryx for JAX debug_info
0156890 
See jax-ml/jax#26480.

PiperOrigin-RevId: 726437113
copybara-service bot pushed a commit that referenced this issue Feb 13, 2025
@gnecula @Google-ML-Automation
[better_errors] Make it explicit that debug_info is not None.
8103ab6 
Now all internal uses of lu.wrap_init and core.Jaxpr are with actual
debug info. This enables us to clean up the type declarations and
to remove the checks whether debug_info is present.

For usage outside of the JAX internals, we change
`jax.extend.linear_util.wrap_init` to be usable without debug_info,
for temporary backwards compatibility. We emit a deprecation
warning and fill-in some fake debugging info.

See #26480 for more details.

PiperOrigin-RevId: 725512692
copybara-service bot pushed a commit that referenced this issue Feb 14, 2025
@gnecula @Google-ML-Automation
[better_errors] Make it explicit that debug_info is not None.
1646061 
Now all internal uses of lu.wrap_init and core.Jaxpr are with actual
debug info. This enables us to clean up the type declarations and
to remove the checks whether debug_info is present.

For usage outside of the JAX internals, we change
`jax.extend.linear_util.wrap_init` to be usable without debug_info,
for temporary backwards compatibility. We emit a deprecation
warning and fill-in some fake debugging info.

See #26480 for more details.

PiperOrigin-RevId: 725512692
copybara-service bot pushed a commit that referenced this issue Feb 14, 2025
@gnecula @Google-ML-Automation
[better_errors] Make it explicit that debug_info is not None.
a0812cd 
Now all internal uses of lu.wrap_init and core.Jaxpr are with actual
debug info. This enables us to clean up the type declarations and
to remove the checks whether debug_info is present.

For usage outside of the JAX internals, we change
`jax.extend.linear_util.wrap_init` to be usable without debug_info,
for temporary backwards compatibility. We emit a deprecation
warning and fill-in some fake debugging info.

See #26480 for more details.

PiperOrigin-RevId: 726770483
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Assignees

@gnecula gnecula

Labels
better_errors Improve the error reporting
Projects
None yet
Milestone
No milestone
Development

No branches or pull requests
1 participant
@gnecula