Commits

Breadcrumbs

History for

on39e4f7f

User selector

All users

Datepicker

All time

Commit History

Commits on Dec 12, 2024

• [sharding_in_types] Make jnp.where broadcast shardings properly when a scalar exists

  Show description for 39e4f7f
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  39e4f7f

  Copy full SHA for 39e4f7f

Commits on Dec 11, 2024

• Use boolean values for partial mask blocks in the splash attention kernel.

  Show description for e92ca9b
  

  Rifur13

  authored and

  Google-ML-Automation

  committed
  e92ca9b

  Copy full SHA for e92ca9b

• jax.lax: raise TypeError for mismatched dtypes

  jakevdp

  committed
  65d2ca6

  Copy full SHA for 65d2ca6

• Increase shard count after adding more tests

  Show description for 20236f1
  

  Rifur13

  authored and

  Google-ML-Automation

  committed
  20236f1

  Copy full SHA for 20236f1

• Merge pull request jax-ml#25395 from gnecula:poly_better_eq

  Show description for 01206f8
  Google-ML-Automation

  committed
  01206f8

  Copy full SHA for 01206f8

• [Mosaic GPU] Split layout inference and dialect lowering files and tests.

  Show description for 4ef7706
  

  bchetioui

  authored and

  Google-ML-Automation

  committed
  4ef7706

  Copy full SHA for 4ef7706

• [Mosaic GPU] Add an initial skeleton for a layout inference pass.

  Show description for 07a3515
  

  bchetioui

  authored and

  Google-ML-Automation

  committed
  07a3515

  Copy full SHA for 07a3515

• Merge pull request jax-ml#25390 from jakevdp:matvec

  Show description for b79dae8
  Google-ML-Automation

  committed
  b79dae8

  Copy full SHA for b79dae8

• [Pallas] Remove grid=1 in tests

  Show description for 13ce517
  

  ayaka14732

  authored and

  Google-ML-Automation

  committed
  13ce517

  Copy full SHA for 13ce517

• Merge pull request jax-ml#24805 from andportnoy:aportnoy/mosaic-gpu-cupti-profiler

  Show description for 0d7eaeb
  Google-ML-Automation

  committed
  0d7eaeb

  Copy full SHA for 0d7eaeb

• Activate Triangular Solve to XLA's FFI

  Show description for 1256153
  

  Paweł Paruzel

  authored and

  Google-ML-Automation

  committed
  1256153

  Copy full SHA for 1256153

• [Mosaic GPU] Automatically format the Mosaic GPU dialect test python code

  Show description for 3d9c720
  

  dimitar-asenov

  authored and

  Google-ML-Automation

  committed
  3d9c720

  Copy full SHA for 3d9c720

• [shape_poly] Improve reasoning for >= in presence of == constraints.

  Show description for 60f9da5
  gnecula

  committed
  60f9da5

  Copy full SHA for 60f9da5

• [Mosaic GPU] Add WGMMA to the Mosaic GPU MLIR Dialect.

  Show description for 66f45d0
  

  dimitar-asenov

  authored and

  Google-ML-Automation

  committed
  66f45d0

  Copy full SHA for 66f45d0

• Merge pull request jax-ml#25392 from carlosgmartin:add_nn_relu_grad_at_zero_test_update_paper_link

  Show description for cfdac00
  Google-ML-Automation

  committed
  cfdac00

  Copy full SHA for cfdac00

• [sharding_in_types] Default axis_types to Auto for all axis_names if user does not set any AxisType. Also resolve some TODOs now that we have a way for user to set the mesh.

  Show description for 41f490a
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  41f490a

  Copy full SHA for 41f490a

• [sharding_in_types] Enforce AxisTypes to always exist if set_mesh is used.

  Show description for b5e4fd1
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  b5e4fd1

  Copy full SHA for b5e4fd1

• Add test of relu grad at zero. Update paper links.

  carlosgmartin

  committed
  0880114

  Copy full SHA for 0880114

• [Pallas TPU] Add WeirdOp to TPU dialect and add lowering for lax.is_finite

  Show description for e88b578
  

  ayaka14732

  authored and

  Google-ML-Automation

  committed
  e88b578

  Copy full SHA for e88b578

• Merge pull request jax-ml#25361 from Rifur13:regression

  Show description for 3ca9f14
  Google-ML-Automation

  committed
  3ca9f14

  Copy full SHA for 3ca9f14

• jax.numpy: implement matvec & vecmat

  jakevdp

  committed
  f6d5876

  Copy full SHA for f6d5876

Commits on Dec 10, 2024

• Only run tpu_all_gather_test on tpu_v5e_4x2

  Show description for 1c1a17e
  

  jburnim

  authored and

  Google-ML-Automation

  committed
  1c1a17e

  Copy full SHA for 1c1a17e

• [jax:custom_partitioning] Make SdyShardingRule a user facing class.

  Show description for d4899f7
  

  bixia1

  authored and

  Google-ML-Automation

  committed
  d4899f7

  Copy full SHA for d4899f7

• Merge pull request jax-ml#25357 from jakevdp:core-deps

  Show description for 9c32fe8
  Google-ML-Automation

  committed
  9c32fe8

  Copy full SHA for 9c32fe8

• [Pallas] Add non-square pl.dot test cases.

  Show description for e418e88
  

  WindQAQ

  authored and

  Google-ML-Automation

  committed
  e418e88

  Copy full SHA for e418e88

• jax.core: deprecate a number of APIs

  jakevdp

  committed
  6541a62

  Copy full SHA for 6541a62

• Reverts bdadc53

  Show description for cb6881d
  

  beckerhe

  authored and

  Google-ML-Automation

  committed
  cb6881d

  Copy full SHA for cb6881d

• Merge pull request jax-ml#25369 from jax-ml:mutable-arrays-ad

  Show description for 263d4d1
  Google-ML-Automation

  committed
  263d4d1

  Copy full SHA for 263d4d1

• Introduce lax.ragged_all_to_all primitive

  Show description for 12c3057
  

  ghpvnist

  authored and

  Google-ML-Automation

  committed
  12c3057

  Copy full SHA for 12c3057

• Add a no-op batching rule for optimization_barrier_p

  Show description for 944d822
  

  yashk2810

  authored and

  Google-ML-Automation

  committed
  944d822

  Copy full SHA for 944d822

• Add a freeze primitive to delimit ref lifetimes for AD.

  Show description for fc2edbf
  

  dougalm

  and

  mattjj

  committed
  fc2edbf

  Copy full SHA for fc2edbf

Commits on Dec 9, 2024

• Fix expected exception type in pallas grad tests.

  Show description for 978d35f
  

  dfm

  authored and

  Google-ML-Automation

  committed
  978d35f

  Copy full SHA for 978d35f

• Adding more tests for multi-head attention

  Rifur13

  committed
  e1e174f

  Copy full SHA for e1e174f

• Disable pjit ArrayPjitTest.test_device_put_grad test on TPU v5e

  Show description for 66b9005
  

  kanglant

  authored and

  Google-ML-Automation

  committed
  66b9005

  Copy full SHA for 66b9005

• Merge pull request jax-ml#25355 from IvyZX:loop-fix

  Show description for 4db533b
  Google-ML-Automation

  committed
  4db533b

  Copy full SHA for 4db533b

Pagination

Previous

Next