mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-18 12:56:07 +00:00
Code Issues Packages Projects Releases Wiki Activity
22,750 Commits 504 Branches 263 Tags
Commit Graph

903 Commits

Author SHA1 Message Date
Jake VanderPlas
f2ffe7f8f2 Deprecate jax.numpy.round_ 
NumPy removed np.round in version 2.0; jax.numpy.round is drop-in
replacement.
 2024-09-03 06:52:07 -07:00
Fabian Pedregosa
530ed026b8 FIX typo on jax.numpy.where docstring 
this was preventing the link to be correctly rendered in the webpage https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.where.html

PiperOrigin-RevId: 670385290
 2024-09-02 20:50:00 -07:00
rajasekharporeddy
ab0e84b4e6 Better docs for jax.numpy: round, around and round_ 2024-09-02 20:39:05 +05:30
Jake VanderPlas
e691f1f36b Fix inaccuracy in jnp.matmul docs 2024-08-29 17:45:43 -07:00
selamw1
76583c87bc gcd_lcm_docstring_added 
description_improved
 2024-08-27 15:39:39 -07:00
jax authors
6ee41369ee Merge pull request #23233 from rajasekharporeddy:testbranch1 
PiperOrigin-RevId: 668061985
 2024-08-27 11:16:11 -07:00
rajasekharporeddy
2f3d428e78 Improved docs for jnp.fix and jnp.trunc 2024-08-27 19:40:28 +05:30
Jake VanderPlas
416f79bb5c DOC: update docstrings for broadcast-related functions 2024-08-26 10:48:29 -07:00
jax authors
c6c701e6a7 Merge pull request #23196 from jakevdp:register-deprecations 
PiperOrigin-RevId: 666900363
 2024-08-23 13:16:09 -07:00
selamw1
ef8532bff5 roll_docstring_added 
roll_docstring_added

see_also_doc_fixed

examples_adjusted
 2024-08-22 10:51:30 -07:00
Jake VanderPlas
2c221f2d5a Register several jax.numpy argument name deprecations 2024-08-22 09:41:53 -07:00
jax authors
9556481be5 Merge pull request #23119 from jakevdp:copy-doc 
PiperOrigin-RevId: 665418823
 2024-08-20 10:21:28 -07:00
jax authors
ae211a135d Merge pull request #23102 from rajasekharporeddy:testbranch1 
PiperOrigin-RevId: 664903092
 2024-08-19 11:38:23 -07:00
jax authors
d115c1cdf1 Merge pull request #23045 from selamw1:diag_docstring 
PiperOrigin-RevId: 664867318
 2024-08-19 10:27:51 -07:00
Jake VanderPlas
a213d2fa30 Improve documentation for jnp.copy 2024-08-19 10:00:37 -07:00
rajasekharporeddy
351b780c48 Better docs for jnp.fmin and fmax 2024-08-19 22:25:53 +05:30
rajasekharporeddy
e39c8685d0 Remove the unused _CONV_PREFERRED_ELEMENT_TYPE_DESCRIPTION 2024-08-17 22:39:10 +05:30
John QiangZhang
ce4dae1c9c Add logging the jax2tf mlir_module_serialized module size. 
PiperOrigin-RevId: 662574156

diag_docstring_added

diag_docstring_desc_fixed

white_space_fixed

white_space_fixed

white_space_fixed

linting fixed
 2024-08-16 20:11:09 -07:00
Jake VanderPlas
8bd84913a6 Better docs for array, asarray, linspace 
This allows removal of extra_params handling from util.implements
 2024-08-09 13:34:50 -07:00
Jake VanderPlas
4ca341701f Improve documentation for jnp.piecewise & jnp.select 2024-08-08 05:53:03 -07:00
Jake VanderPlas
b45f0fe50f Support empty boolean indexing 2024-08-06 09:56:03 -07:00
Philippe Hamel
ecf9f64240 Nit. Fix missing backticks documentation for jnp.where. 
PiperOrigin-RevId: 659549362
 2024-08-05 07:37:41 -07:00
Jake VanderPlas
4d637c8feb Improve documentation for jnp.from_dlpack 2024-08-03 05:48:48 -07:00
George Necula
ffd2b00516 Add concretization error check in core.min_dim and core.max_dim 
Fixes: #22751
 2024-08-01 07:27:35 +02:00
jax authors
f49f2da8fe Merge pull request #22802 from selamw1:diag_indices_docstring 
PiperOrigin-RevId: 658196557
 2024-07-31 17:21:28 -07:00
selamw1
a11ddfd4bc diag_indices_docstring_added 
see_also_diagonal_added
 2024-07-31 16:14:00 -07:00
Jake VanderPlas
a9690adf7b jnp.where: explicitly use mode='drop' when out-of-bound indices expected 2024-07-31 14:55:56 -07:00
Jake VanderPlas
9a88ecb244 Improve error when indexing with too many indices 2024-07-31 13:57:48 -07:00
Jake VanderPlas
5198db9fdb jnp.repeat: add copy argument for Array API 2024-07-30 14:07:08 -07:00
Jake VanderPlas
8bcd288621 Raise ValueError for complex inputs to jnp.clip and jnp.hypot. 
Such inputs were deprecated in JAX v0.4.27, and have been raising a DeprecationWarning for the last several releases.

PiperOrigin-RevId: 657717875
 2024-07-30 13:49:37 -07:00
Jake VanderPlas
cfa1e78549 Improve documentation for jnp.put, jnp.place, jnp.fill_diagonal 
These are all the APIs that have an inplace parameter
 2024-07-29 08:31:54 -07:00
jax authors
aeff5b61a9 Merge pull request #22080 from vfdev-5:add-device-kwarg-linspace-array 
PiperOrigin-RevId: 656467191
 2024-07-26 11:18:24 -07:00
Sergei Lebedev
8d33a6c9a6 Bumped jaxlib version mypy uses on the CI 
I also enabled unnecessary cast checking, because turns out we have quite
a few of those.
 2024-07-26 11:22:39 +01:00
vfdev-5
76d61f9d8f Added device kwargs to jnp.linspace, jnp.array, jnp.asarray 2024-07-26 00:36:34 +02:00
Peter Hawkins
1ac2085417 Fix "unhashable type" error when passing a jax array as the "repeats" argument to jnp.repeat(). 
PiperOrigin-RevId: 656112851
 2024-07-25 15:22:59 -07:00
jax authors
3ed9acba3a Merge pull request #22669 from hawkinsp:repeat 
PiperOrigin-RevId: 656075132
 2024-07-25 13:45:25 -07:00
Peter Hawkins
f07e963bf0 Simplify jaxpr for jnp.repeat in scalar repeat case. 
Before:
```
In [2]: jax.make_jaxpr(lambda x: jnp.repeat(x, 3, axis=-1))(jnp.arange(12).reshape(3, 4))
Out[2]:
{ lambda ; a:i32[3,4]. let
    b:i32[3,4,1] = broadcast_in_dim[broadcast_dimensions=(0, 1) shape=(3, 4, 1)] a
    c:i32[1,3,1,4,1,1] = reshape[dimensions=None new_sizes=(1, 3, 1, 4, 1, 1)] b
    d:i32[1,3,1,4,3,1] = broadcast_in_dim[
      broadcast_dimensions=(0, 1, 2, 3, 4, 5)
      shape=(1, 3, 1, 4, 3, 1)
    ] c
    e:i32[3,4,3] = reshape[dimensions=None new_sizes=(3, 4, 3)] d
    f:i32[3,12] = reshape[dimensions=None new_sizes=(3, 12)] e
  in (f,) }
```

After:
```
In [2]: jax.make_jaxpr(lambda x: jnp.repeat(x, 3, axis=-1))(jnp.arange(12).reshape(3, 4))
Out[2]:
{ lambda ; a:i32[3,4]. let
    b:i32[3,4,3] = broadcast_in_dim[broadcast_dimensions=(0, 1) shape=(3, 4, 3)] a
    c:i32[3,12] = reshape[dimensions=None new_sizes=(3, 12)] b
  in (c,) }
```
 2024-07-25 15:50:23 -04:00
jax authors
5d352a8b0c Merge pull request #22665 from rajasekharporeddy:testbranch1 
PiperOrigin-RevId: 656047033
 2024-07-25 12:37:17 -07:00
rajasekharporeddy
d717135564 Better docs for jnp.triu_indices_from and tril_indices_from 2024-07-26 00:31:44 +05:30
Jake VanderPlas
81b9db6b80 [array api] streamline astype device implementation 
When this was first implemented, convert_element_type did not yet
have a sharding argument. Now we can simplify things by using it.
 2024-07-25 10:42:05 -07:00
Jake VanderPlas
31abce1c80 register several deprecations in jax.numpy 
This is in preparation for finalizing these deprecations. They include:
- complex->real casting in jnp.astype
- complex inputs to jnp.clip
- complex inputs to jnp.hypot

PiperOrigin-RevId: 656005670
 2024-07-25 10:40:00 -07:00
rajasekharporeddy
4ff69d3853 Fix jnp.triu_indices_from and jnp.tril_indices_from to emulate NumPy's behavior for arrays other than 2-D 2024-07-25 16:26:14 +05:30
rajasekharporeddy
c8f30875a9 Improved docs for jnp.triu_indices and tril_indices 2024-07-24 18:24:05 +05:30
jax authors
ac4a7b1e2e Merge pull request #22588 from rajasekharporeddy:testbranch3 
PiperOrigin-RevId: 655211612
 2024-07-23 10:18:56 -07:00
rajasekharporeddy
f90d0ee014 Improved docs for jnp.tri, tril and triu 2024-07-23 20:58:07 +05:30
George Necula
459b83cf4a Reverts 093b92be8ed7bd979486614325956e88cc474ff1 
PiperOrigin-RevId: 655114622
 2024-07-23 04:32:56 -07:00
Dan Foreman-Mackey
991187aaa8 Fix dtype canonicalization in jnp.indices. 
`jnp.indices` was hard coded to default to `dtype = np.int32`, but it
should default to the canonicalized `np.int64`.

Fixes https://github.com/google/jax/issues/22501
 2024-07-22 15:02:48 -04:00
Yuxin Wu
cbb3d99279
Fix docstring of jnp.nonzero 2024-07-19 18:39:12 -07:00
George Necula
093b92be8e Reverts 5216719996d4468f750725ef70cef6f97ac45c27 
PiperOrigin-RevId: 653237245
 2024-07-17 08:10:01 -07:00
George Necula
7817b6785b [shape_poly] Expand the support for shape polymorphism for jnp.pad 
Handle several new padding modes: wrap, reflect, symmetric, linear_ramp, maximum.
Not all situations are handled; try to give a clear error for the unsupported
cases.

While implementing this, I needed to add shape polymorphism support
also for jnp.linspace.

And I discovered a bug in the implementation of `divmod(0, b)`.
 2024-07-15 17:04:54 +02:00