mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-17 20:36:05 +00:00
Code Issues Packages Projects Releases Wiki Activity
24,089 Commits 504 Branches 263 Tags
Commit Graph

987 Commits

Author SHA1 Message Date
carlosgmartin
1f114b1cf7 Add numpy.put_along_axis. 2024-11-14 15:23:26 -05:00
Jake VanderPlas
9359916321 jnp.bincount: support boolean inputs 2024-11-11 06:42:23 -08:00
dymil
9763044d27 Fix argmin docstring to not say "maximum" 2024-11-08 11:19:56 -05:00
Jake VanderPlas
97e8a4c8c6 Fix signatures test: new axis argument in trim_zeros 2024-11-01 10:15:31 -07:00
Dougal Maclaurin
48f24b6acb Remove ConcreteArray from JAX. It's easy to do trace-time concretization without it. 
PiperOrigin-RevId: 691929385
 2024-10-31 14:06:54 -07:00
Jake VanderPlas
adf1492843 Add some missing jax.numpy documentation 2024-10-25 13:14:44 -07:00
jax authors
644f881a51 Merge pull request #24490 from hawkinsp:searchsorted 
PiperOrigin-RevId: 689364122
 2024-10-24 06:56:32 -07:00
Peter Hawkins
a7d711513c Perform searchsorted binary search using unsigned intermediate values. 
Midpoint computation for a binary search should be performed unsigned, see https://research.google/blog/extra-extra-read-all-about-it-nearly-all-binary-searches-and-mergesorts-are-broken/

In addition, we can avoid the somewhat verbose floor_divide HLO since we know the values in question are positive.
 2024-10-23 15:11:55 -04:00
Jake VanderPlas
9bf1516abe Improve docs for jnp.block 2024-10-23 11:37:19 -07:00
Jake VanderPlas
148f9d6559 Better docs for jnp.cov & jnp.corrcoef 2024-10-23 10:17:00 -07:00
Jake VanderPlas
d6f4ce1612 Better docs for jnp.unwrap 2024-10-23 07:58:31 -07:00
Jake VanderPlas
9038bb2664 Better documentation for jnp.indices 2024-10-22 16:48:36 -07:00
Yash Katariya
f8a1f02d6b [sharding_in_types][Take 2] Add out_type argument to einsum and dot_general to allow specifying for the output type. Right now, it only accept a NamedSharding but in the future we can allow a polymorphic type of: jax.ShapeDtypeStruct | Sharding | Layout. 
Reverts 0b3f0e11fb0c37342b3c05ad5d53f3435b6ca44c

PiperOrigin-RevId: 688663504
 2024-10-22 13:10:43 -07:00
Jake VanderPlas
48dd153e18 Better docs for jnp.insert 2024-10-22 09:20:48 -07:00
Jake VanderPlas
7e38cbd604 Better docs for jnp.fromfunction 2024-10-22 08:42:22 -07:00
Jake VanderPlas
8800fe2870 Better documentation for jnp.lexsort 2024-10-21 16:33:14 -07:00
jax authors
441aeebb29 Merge pull request #24420 from superbobry:maint-2 
PiperOrigin-RevId: 688271404
 2024-10-21 14:22:43 -07:00
Sergei Lebedev
3ad1985e1a Bumped mypy and ruff versions used by pre-commit 2024-10-21 21:58:41 +01:00
Jake VanderPlas
66971a2869 Fix jnp.diff for boolean inputs 2024-10-21 13:35:13 -07:00
Jake VanderPlas
6467d03925 Make jnp.subtract a ufunc 2024-10-21 10:11:51 -07:00
jax authors
e29b93ff3e Merge pull request #24421 from jakevdp:cross-doc 
PiperOrigin-RevId: 688175417
 2024-10-21 10:01:45 -07:00
Jake VanderPlas
a1140e9246 Better docs for jnp.cross 2024-10-21 05:59:22 -07:00
Jake VanderPlas
0a85ba5f82 Better documentation for jnp.load 2024-10-19 06:20:20 -07:00
Yash Katariya
57a95a77ff [sharding_in_types] Support jnp.array with sharding_in_types. When the input array has a sharding, propagate it through without dropping the sharding. 
PiperOrigin-RevId: 687089357
 2024-10-17 16:51:41 -07:00
selamw1
24b6f50938 tile_docstring_added 2024-10-15 15:34:01 -07:00
Jake VanderPlas
0c307fe706 Better docs for jnp.apply_along_axis & apply_over_axes 2024-10-14 14:48:33 -07:00
Jake VanderPlas
9b9fc5afae Improve documentation for jnp.histogram & related APIs 2024-10-14 06:44:54 -07:00
Jake VanderPlas
2f798902b4 Better documentation for jnp.nan_to_num 2024-10-09 06:36:45 -07:00
Jake VanderPlas
45f0e9ad68 Simplify definition of jnp.isscalar 
The new semantics are to return True for any array-like object with zero dimensions.
Previously we only returned True for zero-dimensional array-like objects with a weak type. This ends up being more confusing/suprising than it needs to be, and the weak type dependence is rarely useful in practice.

PiperOrigin-RevId: 682656411
 2024-10-05 07:12:20 -07:00
Tom Natan
ed5ba633d4 Reverts 6cf09f8c24c67ff650b95d174501fff3cb59db0d 
PiperOrigin-RevId: 682440543
 2024-10-04 13:56:27 -07:00
Jake VanderPlas
0a8f46a6ca Fix some typos in lax_numpy.py 2024-10-04 03:55:19 -07:00
Jake VanderPlas
4495daee11 Better docs for jnp.issubdtype & jnp.result_type 2024-10-02 05:46:16 -07:00
Blake Hechtman
ce21a12a07 [JAX] Make a one hot mode of take along axis. 
PiperOrigin-RevId: 681139055
 2024-10-01 13:16:26 -07:00
jax authors
9ad7e2eb42 Merge pull request #24048 from jakevdp:packbits-doc 
PiperOrigin-RevId: 681093674
 2024-10-01 11:15:19 -07:00
Jake VanderPlas
ae374e0096 Document jnp.packbits & jnp.unpackbits 2024-10-01 05:33:25 -07:00
Jake VanderPlas
36782e8319 jnp.mask_indices: add docs & tests 2024-09-30 15:13:41 -07:00
jax authors
c557db0bd8 Merge pull request #23995 from jakevdp:trapezoid-doc 
PiperOrigin-RevId: 680734292
 2024-09-30 15:10:16 -07:00
Jake VanderPlas
36d6bb9013 Better docs for jnp.gradient 
Also remove skip_params option from util.implements, as this was its last usage.
 2024-09-30 13:07:52 -07:00
Jake VanderPlas
7cc5ea4c7d Update docs for jnp.trapezoid 2024-09-29 10:08:27 -07:00
Jake VanderPlas
c0612576de Better documentation for jnp.choose 2024-09-27 10:35:19 -07:00
jax authors
96cf2b81e6 Merge pull request #23921 from rajasekharporeddy:testbranch4 
PiperOrigin-RevId: 679203931
 2024-09-26 10:32:44 -07:00
rajasekharporeddy
6072f97961 Raise ValueError when axis1==axis2 for jnp.trace 2024-09-26 21:38:14 +05:30
jax authors
c07652fd46 Merge pull request #23927 from jakevdp:pad-doc 
PiperOrigin-RevId: 679142773
 2024-09-26 07:37:03 -07:00
rajasekharporeddy
8ffeb2388a Better doc for jnp.trace 2024-09-26 09:39:18 +05:30
Jake VanderPlas
ad6c3a7f64 Improve docs for jnp.pad 2024-09-25 14:41:13 -07:00
Jake VanderPlas
ee6fd5aeb2 Improve documentation for jnp.interp 2024-09-25 10:47:46 -07:00
Tom Natan
6cf09f8c24 Reverts eff00cc4499cfe3f3f24bafda6c1ecf908232ff3 
PiperOrigin-RevId: 678756266
 2024-09-25 10:33:53 -07:00
Tom Natan
eff00cc449 [JAX] add support for gather/scatter batching dims following the new attributes in stablehlo. 
This change also uses the new batching dims for gather/scatter batching rules, to avoid concatenating the indices with iota.

See https://github.com/openxla/stablehlo/pull/2259

PiperOrigin-RevId: 678649138
 2024-09-25 04:53:11 -07:00
Jake VanderPlas
cc885ff875 Better docs for jnp.meshgrid 2024-09-23 12:43:14 -07:00
rajasekharporeddy
41eccd925d Improve docs for jnp.logspace and jnp.geomspace 2024-09-23 20:09:12 +05:30