PR #6747 [Pallas] Introduce jax_import_guard

[Pallas] Introduce jax_import_guard #6747

alanwaketan merged 7 commits into master from alanwaketan/jax

alanwaketan added backport_2.3

alanwaketan requested a review from

will-cromar 1 year ago

alanwaketan requested a review from

JackCaoG 1 year ago

alanwaketan assigned

alanwaketan 1 year ago

JackCaoG commented on 2024-03-14

will-cromar commented on 2024-03-14

alanwaketan force pushed from 788d6fe8 to bcfd99c0 1 year ago

JackCaoG approved these changes on 2024-03-19

initial commit

ce2130e5

Fix linters

f694d02c

Fix comment

fe143834

Move JAX import into make_kernel_from_pallas and use _init_computatio…

d127022f

Guard JAX import for TPU tests only

e42b9727

Fix linters

a82b5c6e

Fix cpu/gpu ci

0115d0fe

alanwaketan force pushed from fb204286 to 0115d0fe 1 year ago

alanwaketan merged 7cf9f10a into master 1 year ago

Reviewers

JackCaoG

miladm

will-cromar

Assignees

alanwaketan

Labels

backport_2.3

Milestone

No milestone