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

Lower jax.numpy.dot to mixed-precision dot_general #16721

Merged
merged 1 commit into from
Jul 21, 2023
Merged

Lower jax.numpy.dot to mixed-precision dot_general #16721

merged 1 commit into from
Jul 21, 2023

Conversation

jakevdp
Copy link
Collaborator

@jakevdp jakevdp commented Jul 13, 2023

Related to #16703

@jakevdp jakevdp self-assigned this Jul 13, 2023
@jakevdp jakevdp added the pull ready Ready for copybara import and testing label Jul 13, 2023
@jakevdp jakevdp force-pushed the dot-mixed-precision branch 4 times, most recently from 1f6ecdd to 6603eb2 Compare July 14, 2023 16:00
@jakevdp jakevdp requested a review from yashk2810 July 14, 2023 16:03
This was referenced Jul 14, 2023
Merged
Merged
@jakevdp jakevdp force-pushed the dot-mixed-precision branch 8 times, most recently from 5c844ee to 5dbf3a8 Compare July 19, 2023 18:02
@jakevdp jakevdp mentioned this pull request Jul 19, 2023
Closed
@jakevdp jakevdp mentioned this pull request Jul 20, 2023
Closed
@jakevdp
Copy link
Collaborator Author

jakevdp commented Jul 20, 2023

The remaining failure here is due to #16805

@jakevdp jakevdp force-pushed the dot-mixed-precision branch 6 times, most recently from 57ac77a to ed2ee8e Compare July 20, 2023 21:30
@jakevdp jakevdp force-pushed the dot-mixed-precision branch 2 times, most recently from c72f3d2 to 321884f Compare July 21, 2023 16:31
@jakevdp jakevdp added pull ready Ready for copybara import and testing and removed pull ready Ready for copybara import and testing labels Jul 21, 2023
@copybara-service copybara-service bot merged commit c8f4650 into jax-ml:main Jul 21, 2023
7 checks passed
@jakevdp jakevdp deleted the dot-mixed-precision branch July 21, 2023 17:39
dfm added a commit to dfm/jax that referenced this pull request Sep 16, 2024
@dfm
Unconditionally lower jnp.dot to lax.dot_general.
e5aa584 
jax-ml#16721 added a condition to lower
calls to `jnp.dot` with scalar inputs to `lax.mul` instead of
`lax.dot_general`. AFAICT, jax-ml#16826
fixed the issue that this was solving, so this condition should no
longer be necessary. Removing this condition simplifies the addition of
new arguments to `dot` and `dot_general`, including the `algorithm`
parameter that I am currently working on in
jax-ml#23574, so now seemed like a good time
to remove it!
@dfm dfm mentioned this pull request Sep 16, 2024
Merged
dfm added a commit to dfm/jax that referenced this pull request Sep 16, 2024
@dfm
Unconditionally lower jnp.dot to lax.dot_general.
30f0c05 
jax-ml#16721 added a condition to lower
calls to `jnp.dot` with scalar inputs to `lax.mul` instead of
`lax.dot_general`. AFAICT, jax-ml#16826
fixed the issue that this was solving, so this condition should no
longer be necessary. Removing this condition simplifies the addition of
new arguments to `dot` and `dot_general`, including the `algorithm`
parameter that I am currently working on in
jax-ml#23574, so now seemed like a good time
to remove it!
dfm added a commit to dfm/jax that referenced this pull request Sep 16, 2024
@dfm
Unconditionally lower jnp.dot to lax.dot_general.
95d67a5 
jax-ml#16721 added a condition to lower
calls to `jnp.dot` with scalar inputs to `lax.mul` instead of
`lax.dot_general`. AFAICT, jax-ml#16826
fixed the issue that this was solving, so this condition should no
longer be necessary. Removing this condition simplifies the addition of
new arguments to `dot` and `dot_general`, including the `algorithm`
parameter that I am currently working on in
jax-ml#23574, so now seemed like a good time
to remove it!
dfm added a commit to dfm/jax that referenced this pull request Sep 16, 2024
@dfm
Unconditionally lower jnp.dot to lax.dot_general.
2dfd556 
jax-ml#16721 added a condition to lower
calls to `jnp.dot` with scalar inputs to `lax.mul` instead of
`lax.dot_general`. AFAICT, jax-ml#16826
fixed the issue that this was solving, so this condition should no
longer be necessary. Removing this condition simplifies the addition of
new arguments to `dot` and `dot_general`, including the `algorithm`
parameter that I am currently working on in
jax-ml#23574, so now seemed like a good time
to remove it!
dfm added a commit to dfm/jax that referenced this pull request Sep 17, 2024
@dfm
Unconditionally lower jnp.dot to lax.dot_general.
3b0f793 
jax-ml#16721 added a condition to lower
calls to `jnp.dot` with scalar inputs to `lax.mul` instead of
`lax.dot_general`. AFAICT, jax-ml#16826
fixed the issue that this was solving, so this condition should no
longer be necessary. Removing this condition simplifies the addition of
new arguments to `dot` and `dot_general`, including the `algorithm`
parameter that I am currently working on in
jax-ml#23574, so now seemed like a good time
to remove it!
dfm added a commit to dfm/jax that referenced this pull request Sep 18, 2024
@dfm
Unconditionally lower jnp.dot to lax.dot_general.
adb924e 
jax-ml#16721 added a condition to lower
calls to `jnp.dot` with scalar inputs to `lax.mul` instead of
`lax.dot_general`. AFAICT, jax-ml#16826
fixed the issue that this was solving, so this condition should no
longer be necessary. Removing this condition simplifies the addition of
new arguments to `dot` and `dot_general`, including the `algorithm`
parameter that I am currently working on in
jax-ml#23574, so now seemed like a good time
to remove it!
dfm added a commit to dfm/jax that referenced this pull request Sep 18, 2024
@dfm
Unconditionally lower jnp.dot to lax.dot_general.
1cc9661 
jax-ml#16721 added a condition to lower
calls to `jnp.dot` with scalar inputs to `lax.mul` instead of
`lax.dot_general`. AFAICT, jax-ml#16826
fixed the issue that this was solving, so this condition should no
longer be necessary. Removing this condition simplifies the addition of
new arguments to `dot` and `dot_general`, including the `algorithm`
parameter that I am currently working on in
jax-ml#23574, so now seemed like a good time
to remove it!
keshavb96 added a commit to keshavb96/jax that referenced this pull request Sep 18, 2024
@keshavb96
Squashed commit of the following:
9171e66 
commit 093c6e9
Merge: e1a77ee d0cb318
Author: Keshav <[email protected]>
Date:   Wed Sep 18 13:37:44 2024 -0700

    Merge remote-tracking branch 'upstream/main' into disable_remat_pass

commit e1a77ee
Author: Keshav <[email protected]>
Date:   Wed Sep 18 13:35:37 2024 -0700

    minor changes

commit d0cb318
Merge: b51c653 bef36c4
Author: jax authors <[email protected]>
Date:   Wed Sep 18 13:34:11 2024 -0700

    Merge pull request jax-ml#23736 from hawkinsp:changelog

    PiperOrigin-RevId: 676111400

commit b51c653
Merge: dbc03cf 57a4b76
Author: jax authors <[email protected]>
Date:   Wed Sep 18 13:33:05 2024 -0700

    Merge pull request jax-ml#23737 from jakevdp:digitize-doc

    PiperOrigin-RevId: 676111220

commit dbc03cf
Author: Dan Foreman-Mackey <[email protected]>
Date:   Wed Sep 18 12:39:58 2024 -0700

    Re-land jax-ml#23261 with appropriate compatibility checks.

    PiperOrigin-RevId: 676092618

commit b164d67
Merge: cd04d0f 541b3a3
Author: jax authors <[email protected]>
Date:   Wed Sep 18 12:05:03 2024 -0700

    Merge pull request jax-ml#23247 from kaixih:sliding_window_attn

    PiperOrigin-RevId: 676079831

commit 57a4b76
Author: Jake VanderPlas <[email protected]>
Date:   Wed Sep 18 11:59:00 2024 -0700

    Improve documentation for jnp.digitize

commit bef36c4
Author: Peter Hawkins <[email protected]>
Date:   Wed Sep 18 18:57:03 2024 +0000

    Add Python 3.13 wheels to changelog.

commit cd04d0f
Merge: 016c499 c756d9b
Author: jax authors <[email protected]>
Date:   Wed Sep 18 10:00:03 2024 -0700

    Merge pull request jax-ml#23726 from hawkinsp:debug

    PiperOrigin-RevId: 676030839

commit 016c499
Author: Sergei Lebedev <[email protected]>
Date:   Wed Sep 18 09:56:44 2024 -0700

    Removed leftover usages of GPUGridSpec from Pallas Mosaic GPU tests

    PiperOrigin-RevId: 676029854

commit 9dd363d
Author: Luke Baumann <[email protected]>
Date:   Wed Sep 18 09:28:25 2024 -0700

    Export `jax.lib.xla_extension.ifrt_programs`.

    PiperOrigin-RevId: 676020419

commit e27f1e9
Author: jax authors <[email protected]>
Date:   Wed Sep 18 09:03:55 2024 -0700

    Change Python version 3.13.0rc2 to 3.13.0-rc.2.

    The value is taken from [the versions manifest](https://raw.githubusercontent.com/actions/python-versions/main/versions-manifest.json).

    PiperOrigin-RevId: 676012255

commit 442e863
Author: Sergei Lebedev <[email protected]>
Date:   Wed Sep 18 08:56:49 2024 -0700

    Added a missing branch to `mgpu.FragmentedArray.astype`

    Previously, an unsupported cast produced a `NameError` instead.

    PiperOrigin-RevId: 676010161

commit 6236b8f
Merge: 826843a 1cc9661
Author: jax authors <[email protected]>
Date:   Wed Sep 18 08:57:38 2024 -0700

    Merge pull request jax-ml#23667 from dfm:always-lower-jnp-dot-to-dot-general

    PiperOrigin-RevId: 676010154

commit c756d9b
Author: Peter Hawkins <[email protected]>
Date:   Wed Sep 18 15:44:45 2024 +0000

    Fix error in debugger tests that is showing up in CI.

    I'm unsure why this started happening now, but sometimes we get an
    invalid offset for a frame. Be tolerant of that case.

commit 826843a
Merge: c191bbc 922e652
Author: jax authors <[email protected]>
Date:   Wed Sep 18 08:42:39 2024 -0700

    Merge pull request jax-ml#23723 from hawkinsp:setuptools

    PiperOrigin-RevId: 676005613

commit c191bbc
Author: Yash Katariya <[email protected]>
Date:   Wed Sep 18 08:40:30 2024 -0700

    Make `debug.print` work with static args. Fixes: jax-ml#23600

    PiperOrigin-RevId: 676005582

commit 1cc9661
Author: Dan Foreman-Mackey <[email protected]>
Date:   Mon Sep 16 14:18:29 2024 -0400

    Unconditionally lower jnp.dot to lax.dot_general.

    jax-ml#16721 added a condition to lower
    calls to `jnp.dot` with scalar inputs to `lax.mul` instead of
    `lax.dot_general`. AFAICT, jax-ml#16826
    fixed the issue that this was solving, so this condition should no
    longer be necessary. Removing this condition simplifies the addition of
    new arguments to `dot` and `dot_general`, including the `algorithm`
    parameter that I am currently working on in
    jax-ml#23574, so now seemed like a good time
    to remove it!

commit 922e652
Author: Peter Hawkins <[email protected]>
Date:   Wed Sep 18 15:17:49 2024 +0000

    Replace plat-name with plat_name.

    The former seems to elicit a deprecation warning from setuptools
    recently.

commit 69ba060
Author: Dan Foreman-Mackey <[email protected]>
Date:   Wed Sep 18 07:40:58 2024 -0700

    Reverts e15ec1e

    PiperOrigin-RevId: 675987338

commit 44a7f04
Merge: 0a29696 2834c13
Author: jax authors <[email protected]>
Date:   Wed Sep 18 07:31:00 2024 -0700

    Merge pull request jax-ml#23708 from jakevdp:sort-complex

    PiperOrigin-RevId: 675983957

commit 0a29696
Merge: e15ec1e 73c38cb
Author: jax authors <[email protected]>
Date:   Wed Sep 18 07:08:24 2024 -0700

    Merge pull request jax-ml#23698 from dfm:dev-clang-warning

    PiperOrigin-RevId: 675977448

commit 2834c13
Author: Jake VanderPlas <[email protected]>
Date:   Tue Sep 17 15:32:25 2024 -0700

    jnp.sort_complex: fix output for N-dimensional inputs

commit 73c38cb
Author: Dan Foreman-Mackey <[email protected]>
Date:   Tue Sep 17 14:00:21 2024 -0400

    Add a note to the developer docs making it clear that clang is the only
    toolchain that is actively supported for source compilation.

    As discussed in jax-ml#23687

commit e15ec1e
Merge: 48d8fce 3f2bc9b
Author: jax authors <[email protected]>
Date:   Wed Sep 18 06:56:28 2024 -0700

    Merge pull request jax-ml#23261 from joaospinto:stablehlo.tan

    PiperOrigin-RevId: 675973798

commit 48d8fce
Merge: 4e6f690 2714469
Author: jax authors <[email protected]>
Date:   Wed Sep 18 06:54:28 2024 -0700

    Merge pull request jax-ml#23563 from rajasekharporeddy:testbranch1

    PiperOrigin-RevId: 675973225

commit 4e6f690
Merge: b7c91e9 611ad63
Author: jax authors <[email protected]>
Date:   Wed Sep 18 06:35:15 2024 -0700

    Merge pull request jax-ml#23653 from apaszke:torchsaic

    PiperOrigin-RevId: 675967844

commit b7c91e9
Author: Sergei Lebedev <[email protected]>
Date:   Wed Sep 18 06:22:14 2024 -0700

    Lookup `shape` and `dtype` directly on `state.AbstractRef` instead of going through `inner_aval`

    This is just a cleanup. No behavior changes are expected.

    PiperOrigin-RevId: 675964703

commit 611ad63
Author: Adam Paszke <[email protected]>
Date:   Fri Sep 6 16:09:58 2024 +0000

    Add basic PyTorch integration for Mosaic GPU

    We have already had most of the relevant pieces and we only needed
    to connect them together. The most sensitive change is perhaps that
    I needed to expose one more symbol from the XLA GPU plugin, but I don't
    think it should be a problem.

commit e903369
Author: Sergei Lebedev <[email protected]>
Date:   Wed Sep 18 05:25:37 2024 -0700

    Pulled `scratch_shapes` into `GridSpec`

    It is supported by Mosaic TPU and Mosaic GPU and unsupported by Triton.

    PiperOrigin-RevId: 675950199

commit 2714469
Author: rajasekharporeddy <[email protected]>
Date:   Wed Sep 18 17:06:28 2024 +0530

    Deprecate passing NdArrays with ndim != 1 and non-arraylike inputs to jnp.trim_zeros

commit b904599
Author: Sergei Lebedev <[email protected]>
Date:   Wed Sep 18 04:23:25 2024 -0700

    `pl.debug_print` no longer restricts values to be scalars

    This allows printing arrays on Triton and soon on Mosaic GPU.

    PiperOrigin-RevId: 675935666

commit 988ed2b
Author: jax authors <[email protected]>
Date:   Tue Sep 17 21:09:26 2024 -0700

    Add support for SMEM windows in Pallas custom pipeline.

    PiperOrigin-RevId: 675822640

commit f79d85b
Merge: 1b74cfd cc28d63
Author: Keshav <[email protected]>
Date:   Tue Sep 17 18:58:33 2024 -0700

    Merge remote-tracking branch 'upstream/main' into disable_remat_pass

commit cc28d63
Merge: 8bcdb12 9d3762b
Author: jax authors <[email protected]>
Date:   Tue Sep 17 17:36:36 2024 -0700

    Merge pull request jax-ml#23682 from sharadmv:pallas-async-docs

    PiperOrigin-RevId: 675770723

commit 1b74cfd
Author: Keshav <[email protected]>
Date:   Tue Sep 17 17:23:30 2024 -0700

    disable remat hlo pass by default

commit 8bcdb12
Author: jax authors <[email protected]>
Date:   Tue Sep 17 16:50:55 2024 -0700

    Add CI jobs for python 3.13.0rc2.

    PiperOrigin-RevId: 675758096

commit 8b5b717
Author: Yash Katariya <[email protected]>
Date:   Tue Sep 17 16:39:55 2024 -0700

    Fix jaxpr equation context propagation in jaxpr equations when `inline=True`.

    PiperOrigin-RevId: 675754808

commit 86fe463
Author: Parker Schuh <[email protected]>
Date:   Tue Sep 17 16:10:41 2024 -0700

    [Take 2] Generalize global jit cpp cache keys so we can add more keys than the current donate_argnums.

    This allows us to get more cache hits globally. For example:

    Before:

    jax.jit(f, out_shardings=s)(arr)
    jax.jit(f, out_shardings=s)(arr)  # cpp cache miss
    After:

    jax.jit(f, out_shardings=s)(arr)
    jax.jit(f, out_shardings=s)(arr)  # cpp cache hit

    Reverts b615266

    PiperOrigin-RevId: 675746175

commit e92a599
Author: Christos Perivolaropoulos <[email protected]>
Date:   Tue Sep 17 15:26:42 2024 -0700

    [mosaic_gpu] Better error message for misaligned tma_transpose with dtype.

    PiperOrigin-RevId: 675731295

commit 7864648
Merge: 3f2c58b 83a7555
Author: jax authors <[email protected]>
Date:   Tue Sep 17 15:12:50 2024 -0700

    Merge pull request jax-ml#23679 from selamw1:docstring_sort_complex

    PiperOrigin-RevId: 675726527

commit 83a7555
Author: selamw1 <[email protected]>
Date:   Mon Sep 16 16:47:52 2024 -0700

    docstring_sort_complex_added

    input_array_modified

commit 9d3762b
Author: Sharad Vikram <[email protected]>
Date:   Mon Sep 16 19:18:22 2024 -0700

    [Pallas] Add design note for async ops on TPU

commit 3f2bc9b
Author: Joao Sousa-Pinto <[email protected]>
Date:   Mon Aug 26 17:25:16 2024 -0700

    Lower tan to StableHLO instead of CHLO.

    Fixes jax-ml#23259

commit 541b3a3
Author: kaixih <[email protected]>
Date:   Mon Aug 26 17:32:38 2024 +0000

    New feature
rajasekharporeddy pushed a commit to rajasekharporeddy/jax that referenced this pull request Sep 20, 2024
@dfm @rajasekharporeddy
Unconditionally lower jnp.dot to lax.dot_general.
9af26b9 
jax-ml#16721 added a condition to lower
calls to `jnp.dot` with scalar inputs to `lax.mul` instead of
`lax.dot_general`. AFAICT, jax-ml#16826
fixed the issue that this was solving, so this condition should no
longer be necessary. Removing this condition simplifies the addition of
new arguments to `dot` and `dot_general`, including the `algorithm`
parameter that I am currently working on in
jax-ml#23574, so now seemed like a good time
to remove it!
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@yashk2810 yashk2810 yashk2810 approved these changes

Assignees

@jakevdp jakevdp

Labels
pull ready Ready for copybara import and testing
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
3 participants
@jakevdp @yashk2810 @kokoro-team