PR #148 CI: 11/22/24 upstream sync

Add a link to Intel plugin for JAX

mini-goel committed 1 year ago

Add float8_e4m3 and float8_e3m4 types support

apivovarov committed 1 year ago

Move Control Flow text from Sharp Bits into its own tutorial.

emilyfertig committed 1 year ago

jnp.logaddexp2: simplify implementation

jakevdp committed 1 year ago

Update array-api-tests commit

jakevdp committed 1 year ago

Merge pull request #24864 from jakevdp:logaddexp2

Google-ML-Automation committed 1 year ago

Add numpy.put_along_axis.

carlosgmartin committed 1 year ago

Merge pull request #24904 from jakevdp:array-api

Google-ML-Automation committed 1 year ago

cleanup: delete unused argument from internal reduction helper

jakevdp committed 1 year ago

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

Add shard map replication rule for ffi_call.

dfm committed 1 year ago

Merge pull request #24567 from Intel-tensorflow:minigoel/intel-plugin

Google-ML-Automation committed 1 year ago

Merge pull request #24905 from jakevdp:old-arg

Google-ML-Automation committed 1 year ago

Merge pull request #24871 from carlosgmartin:numpy_put_along_axis

Google-ML-Automation committed 1 year ago

Update array-api-tests commit

jakevdp committed 1 year ago

Merge pull request #24881 from dfm:ffi-call-rep-rule

Google-ML-Automation committed 1 year ago

Merge pull request #24907 from jakevdp:array-api

Google-ML-Automation committed 1 year ago

Merge pull request #24862 from emilyfertig:emilyaf-control-flow-tutorial

Google-ML-Automation committed 1 year ago

[sharding_in_types] Handle collective axes in lowering rules more generally. If any axis is collective, set all dims of aval to unspecified dims in `wrap_with_sharding_op`.

Google-ML-Automation committed 1 year ago

Set __module__ attribute for objects in jax.numpy

jakevdp committed 1 year ago

[Mosaic TPU] Support 1D concat: set implicit_dim to kSecondMinor to treat 1D (N,) as (1, N) and then tile it as (1, 128)

Google-ML-Automation committed 1 year ago

Lower threefry as an out-of-line MLIR function on TPU.

hawkinsp committed 1 year ago

Merge pull request #24913 from hawkinsp:threefry

Google-ML-Automation committed 1 year ago

Add missing functions to jax.numpy type interface

jakevdp committed 1 year ago

Add an example on logical operators to the tutorial.

emilyfertig committed 1 year ago

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

DOC: update main landing page style

jakevdp committed 1 year ago

Merge pull request #24918 from emilyfertig:emilyaf-logical-op-example

Google-ML-Automation committed 1 year ago

Merge pull request #24914 from jakevdp:fix-pyi

Google-ML-Automation committed 1 year ago

Consolidate material on PRNGs and add a short summary to Key Concepts.

emilyfertig committed 1 year ago

Merge pull request #24917 from emilyfertig:emilyaf-sharp-bits

Google-ML-Automation committed 1 year ago

[sharding_in_types] Don't emit a wsc under full manual mode to avoid increasing HLO size by a lot

Google-ML-Automation committed 1 year ago

Adds a flag to control proxy env checking.

yliu120 committed 1 year ago

Deduplicate constants in StableHLO lowering.

hawkinsp committed 1 year ago

Return a ndarray in shape_as_value if the shape is known to be constant.

hawkinsp committed 1 year ago

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

Use a direct StableHLO lowering for pow.

hawkinsp committed 1 year ago

Adds an env that can let users provide a custom version suffix for jax dev build.

yliu120 committed 1 year ago

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

[AutoPGLE] Temporary disable pgle_test in the OSS.

Google-ML-Automation committed 1 year ago

Merge pull request #24930 from hawkinsp:dedup

Google-ML-Automation committed 1 year ago

Merge pull request #24933 from hawkinsp:pow

Google-ML-Automation committed 1 year ago

Add a GPU implementation of `lax.linalg.eig`.

Google-ML-Automation committed 1 year ago

Merge pull request #24932 from hawkinsp:gather

Google-ML-Automation committed 1 year ago

Add new CI script for running Bazel GPU presubmits

Google-ML-Automation committed 1 year ago

Merge pull request #24912 from jakevdp:jnp-module

Google-ML-Automation committed 1 year ago

Merge pull request #24916 from jakevdp:update-lp

Google-ML-Automation committed 1 year ago

Make logaddexp and logaddexp2 into ufuncs

jakevdp committed 1 year ago

Merge pull request #24903 from jakevdp:logsumexp

Google-ML-Automation committed 1 year ago

Return SingleDeviceSharding instead of GSPMDShardings when there is only 1 device during `compiled.input_shardings` call.

Google-ML-Automation committed 1 year ago

fix typo in numpy/__init__.pyi

jakevdp committed 1 year ago

Merge pull request #24952 from jakevdp:fix-pyi

Google-ML-Automation committed 1 year ago

[SDY] fix JAX layouts tests for Shardy.

Google-ML-Automation committed 1 year ago

Disable some complex function accuracy tests that fail on Mac ARM.

hawkinsp committed 1 year ago

Filter custom dtypes by supported_dtypes in `_LazyDtypes`.

Google-ML-Automation committed 1 year ago

[Pallas] Increase test coverage of pl.dot.

Google-ML-Automation committed 1 year ago

Merge pull request #24957 from hawkinsp:arm

Google-ML-Automation committed 1 year ago

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

Merge pull request #23585 from apivovarov:float8_e4m3

Google-ML-Automation committed 1 year ago

Adds font fallbacks

barnesjoseph committed 1 year ago

Merge pull request #24958 from barnesjoseph:add-font-fallback

Google-ML-Automation committed 1 year ago

Delete _normalized_spec from NamedSharding

Google-ML-Automation committed 1 year ago

Fix a bug where mesh checking was not correct

Google-ML-Automation committed 1 year ago

[pallas] Minor simplifications to Pallas interpreter.

Google-ML-Automation committed 1 year ago

Update jax.scipy.special.gamma and gammasgn to return NaN for negative integer inputs.

hawkinsp committed 1 year ago

Merge pull request #24945 from hawkinsp:gamma

Google-ML-Automation committed 1 year ago

Merge pull request #24861 from yliu120:add_versions

Google-ML-Automation committed 1 year ago

[Mosaic TPU] Support relayout for mask vector

Google-ML-Automation committed 1 year ago

Merge pull request #24853 from yliu120:check_proxy_envs

Google-ML-Automation committed 1 year ago

Implement lax.pad in Pallas.

Google-ML-Automation committed 1 year ago

[AutoPGLE] Use compile options to override debug options instead of XLA_FLAGS.

Google-ML-Automation committed 1 year ago

Fix some typos

nireekshak committed 1 year ago

fix(docs): typos in macro name

jeertmans committed 1 year ago

Merge pull request #24942 from jeertmans:patch-1

Google-ML-Automation committed 1 year ago

Merge pull request #24968 from nireekshak:testingbranch

Google-ML-Automation committed 1 year ago

Add missing version guard in GPU tests for jnp.poly.

Google-ML-Automation committed 1 year ago

[Mosaic TPU] Add general tpu.vector_store and support masked store.

Google-ML-Automation committed 1 year ago

Add alternate implementation of threefry as a pallas kernel.

Google-ML-Automation committed 1 year ago

Add test utility for accessing jaxlib version tuple.

Google-ML-Automation committed 1 year ago

Add a new API jax.lax.split.

hawkinsp committed 1 year ago

Merge pull request #24970 from hawkinsp:split

Google-ML-Automation committed 1 year ago

[pallas:mosaic_gpu] `copy_gmem_to_smem` no longer requires `barrier` to be a keyword argument

Google-ML-Automation committed 1 year ago

[Mosaic] Add target core type parameter to tpu.sem_signal

Google-ML-Automation committed 1 year ago

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

Move JAX example to public XLA:CPU API

Google-ML-Automation committed 1 year ago

Reverts 2075b091c4e83f0bdbd0d47812a72114fb8b937a

Google-ML-Automation committed 1 year ago

[Mosaic] Extend tpu.sem_signal with subcore_id

Google-ML-Automation committed 1 year ago

Make deprecated jax.experimental.array_api module visibility internal-only

Google-ML-Automation committed 1 year ago

Fix a bug where constant deduplication used an inappropriate inequality.

Google-ML-Automation committed 1 year ago

[Mosaic] Add `tpu.log` verification on SC

Google-ML-Automation committed 1 year ago

represent `random.key_impl` of builtin RNGs by canonical string name

froystig committed 1 year ago

Add test_compute_on_host_shared_sharding in memories_test

Google-ML-Automation committed 1 year ago

Merge pull request #24593 from froystig:random-dtypes

Google-ML-Automation committed 1 year ago

[mosaic_gpu] Fix signedness handling in FragmentedArray._pointwise.

Google-ML-Automation committed 1 year ago

[pallas:mosaic_gpu] Add Pallas Mosaic GPU lowering for jax.lax.bitcast_convert_type.

Google-ML-Automation committed 1 year ago

[pallas:mosaic_gpu] `copy_smem_to_gmem` now supports `wait_read_only`

Google-ML-Automation committed 1 year ago

[mosaic_gpu] Fixed `FragmentedArray` comparisons with literals

Google-ML-Automation committed 1 year ago

[mosaic_gpu] Handle older `jaxlib`s in the profiler module

Google-ML-Automation committed 1 year ago

[pallas] Do not skip vmap tests on GPU when x64 is enabled

Google-ML-Automation committed 1 year ago

Update XLA dependency to use revision

Google-ML-Automation committed 1 year ago

jax CI: 11/22/24 upstream sync #148 Merged

jax
CI: 11/22/24 upstream sync
#148

Merged