ROCm/jax

Pull Requests Commits

Added transpose rhs option to ragged dot kernel, useful for implementing the backward function.

Google-ML-Automation committed 304 days ago

d8f9a27a

Merge pull request #33017 from mattjj:vjp3-slicing-fancy-transposes

Google-ML-Automation committed 304 days ago

95468afe

[Pallas TPU] Automatically plumb name argument to pallas_call into name_stack

sharadmv committed 304 days ago

b7d74505

Merge pull request #33015 from mattjj:shmap-varargs-error-2

Google-ML-Automation committed 304 days ago

7353de97

[mutable-arrays] add fancy transpose rules for slice / dynamic_slice

mattjj committed 304 days ago

560c1aa8

Add additional way to poision a connection (to allow testing different

pschuh committed 304 days ago

88a8b033

[JAX] Import MLIR dialects from jax via jax.extend.mlir, not jax._src.lib.mlir.

hawkinsp committed 304 days ago

03614c96

Catch overflow if very large 'k' are passed to jnp.tri.

hawkinsp committed 304 days ago

075494cd

Merge pull request #32990 from arvindsankar:ensure_arraylike

Google-ML-Automation committed 304 days ago

dea87084

Add missing debug_info args in physicalize_jaxpr.

jburnim committed 304 days ago

eca5b593

[NFC] Add `name` argument to `pallas_call` primitive.

sharadmv committed 304 days ago

473499aa

[shmap] fix error messages with varargs and kwargs

mattjj committed 304 days ago

555d502d

Make MetaTy work with HiJAX. Before this HiTypes (with values) were being held in Traced without converting it to MetaTys.

yashk2810 committed 304 days ago

7d256f84

nn: use ensure_arraylike instead of check_arraylike and jnp.asarray

arvindsankar committed 304 days ago

5e72730f

When we do `traced = jit(f).trace(duck_stuff)`, we only use duck_stuff to infer the input type (and other metadata like layouts) at which to trace f. Because duck_stuff is flexible, it could comprise actual values, but we don't need to hold onto those values. Yet `traced` currently does hold references to any such values!

yashk2810 committed 304 days ago

d2ce04b6

Add dynamic grid with smem output test.

WindQAQ committed 304 days ago

619dbe12

Update Bazel version to 7.7.0.

meteorcloudy committed 304 days ago

78fce199

Remove unnecessary static dimension check on the k argument to tri.

hawkinsp committed 304 days ago

6671ddbe

Use significantly fewer processes for Windows nightlies under pytest.

belitskiy committed 304 days ago

c0d8bdbd

Remove/fix a couple of old TODOs.

hawkinsp committed 304 days ago

b82ed8ff

[Pallas:MGPU] Make `TransposeRef` a frozen dataclass.

allanrenucci committed 304 days ago

eace2086

Set __module__ on jax.ShapeDtypeStruct.

hawkinsp committed 304 days ago

e62f7243

[pallas:mosaic] Inlined `tpu_info.supports_sparse_core`

superbobry committed 304 days ago

0772c779

Update XLA dependency to use revision http://github.com/openxla/xla/commit/5aff798706c762c9ff125092cb8021ccddd31b6b

Google-ML-Automation committed 304 days ago

415e8195

[Mosaic:TPU][Relayouts] Replicated minor large-to-small retiling (one-to-many)

tlongeri committed 305 days ago

bd0397dd

Use copy instead of reuse semantics for CopyToHost, because reuse doesn't make sense in that context.

Google-ML-Automation committed 305 days ago

e7c5fd3f

Merge pull request #32668 from arvindsankar:jnp-mean

Google-ML-Automation committed 305 days ago

55b5aef1

[Mosaic GPU][NFC] Always use pre-order traversal.

allanrenucci committed 305 days ago

3ea50d5d

Remove dead code

yashk2810 committed 305 days ago

533ee6c3

Use TestCase.enterContext in JaxTestCase.setUp

hawkinsp committed 305 days ago

a3a460f5

Newer Older