mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-17 04:16:07 +00:00
Code Issues Packages Projects Releases Wiki Activity
rocm_jax/jax/_src
History
George Necula e643f98558 [shape_poly] Reimplement the shape constraint checking using shape assertions. 
Most of the functionality is for the JAX native serialization case.
This relies on newly added functionality to xla_extension.refine_polymorphic_shapes
that handles custom calls @static_assertion.

As a beneficial side-effect now we get shape constraint checking for jax2tf
graph serialization when the resulting function is executed in graph mode.
2023-07-19 09:56:33 +03:00
..
clusters
Use lower-case PEP 585 names for types.
2023-06-23 15:12:14 -07:00
debugger
Use lower-case PEP 585 names for types.
2023-06-23 15:12:14 -07:00
image
[shape_poly] Cleanups for the shape polymorphism APIs.
2023-06-30 15:56:57 +03:00
internal_test_util
Change to simpler import for jax.config
2023-04-21 11:51:22 -07:00
interpreters
[shape_poly] Reimplement the shape constraint checking using shape assertions.
2023-07-19 09:56:33 +03:00
lax
Merge pull request #16718 from mattjj:scatter-apply-autodiff
2023-07-14 09:24:12 -07:00
lib
Use the upstream MLIR strip-debuginfo pass instead of hand-rolling our own.
2023-07-14 12:24:59 -07:00
nn
[shape_poly] Cleanups for the shape polymorphism APIs.
2023-06-30 15:56:57 +03:00
numpy
support np.array(x) where x is a custom pytree with __jax_array__
2023-07-17 13:33:17 -07:00
ops
Use lower-case PEP 585 names for types.
2023-06-23 15:12:14 -07:00
scipy
Add jax.scipy.stats.binom
2023-06-27 03:41:38 -07:00
state
Add discharge rules for scan/while
2023-07-06 22:30:35 +00:00
third_party
Use lower-case PEP 585 names for types.
2023-06-23 15:12:14 -07:00
__init__.py
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
abstract_arrays.py
Use lower-case PEP 585 names for types.
2023-06-23 15:12:14 -07:00
ad_checkpoint.py
Use lower-case PEP 585 names for types.
2023-06-23 15:12:14 -07:00
ad_util.py
Use lower-case PEP 585 names for types.
2023-06-23 15:12:14 -07:00
api_util.py
jax.jit now works correctly if both donate_argnums and donate_argnames are specified.
2023-07-14 14:28:16 -07:00
api.py
jax.jit now works correctly if both donate_argnums and donate_argnames are specified.
2023-07-14 14:28:16 -07:00
array.py
Use lower-case PEP 585 names for types.
2023-06-23 15:12:14 -07:00
basearray.py
Split basearray into separate Bazel module.
2023-03-13 11:14:41 -07:00
basearray.pyi
Use lower-case PEP 585 names for types.
2023-06-23 15:12:14 -07:00
callback.py
Add maximal sharding for pure_callback not inside of a shard_map
2023-05-11 13:28:37 -07:00
checkify.py
Checkify: remove duplicate line.
2023-07-10 15:41:35 +01:00
cloud_tpu_init.py
Issue warning if JAX_USE_PJRT_C_API_ON_TPU=false
2023-06-20 14:27:08 -07:00
compilation_cache_interface.py
Move implementation of compilation cache out of jax/experimental and into jax/_src.
2023-04-17 08:35:53 -07:00
compilation_cache.py
Use the upstream MLIR strip-debuginfo pass instead of hand-rolling our own.
2023-07-14 12:24:59 -07:00
config.py
Remove the jax_array config option, which does nothing.
2023-07-18 06:16:06 -07:00
core.py
[shape_poly] More cleanup for the internal APIs for shape polymorphism.
2023-07-13 16:37:53 +03:00
custom_api_util.py
Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".
2022-09-22 12:27:19 -07:00
custom_batching.py
Fix type failures under an upcoming pytype change.
2023-04-07 07:10:09 -07:00
custom_derivatives.py
document and test CustomVJPPrimal type as API symbol
2023-07-12 09:19:02 -07:00
custom_transpose.py
Use lower-case PEP 585 names for types.
2023-06-23 15:12:14 -07:00
debugging.py
Fix docstring for jax.debug.{print,callback}
2023-07-17 15:58:58 +01:00
deprecations.py
Make deprecation warnings warn at appropriate stacklevel
2023-05-23 14:43:38 -07:00
dispatch.py
Use lower-case PEP 585 names for types.
2023-06-23 15:12:14 -07:00
distributed.py
[PJRT PLUGIN] Add num_processes to distributed.global_state.
2023-05-26 13:14:40 -07:00
dlpack.py
[JAX] Delete _DeviceArray and DeviceArray.
2023-03-29 15:07:14 -07:00
dtypes.py
type promotion: improve implicit type promotion error for float8 & int4
2023-07-13 10:41:49 -07:00
effects.py
Use lower-case PEP 585 names for types.
2023-06-23 15:12:14 -07:00
environment_info.py
Split _src files custom_api_util, deprecations, effects and environment_info into separate Bazel targets.
2023-03-10 12:26:05 -08:00
errors.py
errors: create TracerBoolConversionError for more targeted debugging tips
2023-06-21 01:41:45 -07:00
flatten_util.py
Avoid imports from the public jax.* namespace in more places internally.
2023-04-04 11:41:40 -07:00
gfile_cache.py
Fix FileExistsError for tests/gfile_cache_test.py::FileSystemCacheTest::test_threads
2023-06-14 08:07:22 +08:00
lax_reference.py
Add padding option "SAME_LOWER" for ticket https://github.com/google/jax/pull/14990
2023-03-23 15:50:16 -07:00
lazy_loader.py
Use lower-case PEP 585 names for types.
2023-06-23 15:12:14 -07:00
linear_util.py
Use lower-case PEP 585 names for types.
2023-06-23 15:12:14 -07:00
maps.py
jax.jit now works correctly if both donate_argnums and donate_argnames are specified.
2023-07-14 14:28:16 -07:00
mesh.py
Use lower-case PEP 585 names for types.
2023-06-23 15:12:14 -07:00
monitoring.py
Use lower-case PEP 585 names for types.
2023-06-23 15:12:14 -07:00
op_shardings.py
Use lower-case PEP 585 names for types.
2023-06-23 15:12:14 -07:00
partition_spec.py
Move PartitionSpec into its own file (jax/_src/partition_spec.py).
2023-04-06 11:43:25 -07:00
path.py
Make the etils[epath] dependency of JAX optional.
2022-10-12 21:16:48 +00:00
pickle_util.py
Use cloudpickle for serializing host callbacks
2023-06-09 13:34:43 -07:00
pjit.py
Delete in_axis_resources and out_axis_resources from pjit since it's been more than 3 months since their deprecation. The replace is to use in_shardings and out_shardings. You can still pass PartitionSpecs to {in|out}_shardings to pjit.
2023-07-17 06:35:49 -07:00
pretty_printer.py
Use lower-case PEP 585 names for types.
2023-06-23 15:12:14 -07:00
prng.py
Remove obsolete jaxlib version checks
2023-07-12 11:53:55 -07:00
profiler.py
Split _src/profiler into a separate BUILD target.
2023-03-10 11:38:46 -08:00
public_test_util.py
Require ml_dtypes>=0.2
2023-07-07 12:07:44 -07:00
random.py
accept general shape option in jax.random.split
2023-07-11 14:24:01 -07:00
sharding_impls.py
Add memory_kind to NamedSharding, SingleDeviceSharding, PositionalSharding and GSPMDSharding.
2023-07-18 07:39:36 -07:00
sharding_specs.py
Use lower-case PEP 585 names for types.
2023-06-23 15:12:14 -07:00
sharding.py
Add memory_kind to NamedSharding, SingleDeviceSharding, PositionalSharding and GSPMDSharding.
2023-07-18 07:39:36 -07:00
source_info_util.py
Use lower-case PEP 585 names for types.
2023-06-23 15:12:14 -07:00
stages.py
Remove deprecated compiler_ir from Compiled
2023-07-11 09:24:48 -07:00
test_util.py
[shape_poly] Reimplement the shape constraint checking using shape assertions.
2023-07-19 09:56:33 +03:00
traceback_util.py
Use lower-case PEP 585 names for types.
2023-06-23 15:12:14 -07:00
tree_util.py
Use lower-case PEP 585 names for types.
2023-06-23 15:12:14 -07:00
typing.py
Allow duck-typed inputs to zeros_like, ones_like, etc.
2023-06-06 00:59:51 -07:00
util.py
Remove obsolete jaxlib version checks
2023-07-12 11:53:55 -07:00
xla_bridge.py
Enable to set fdo_profile through XLA python client.
2023-07-11 14:47:39 -07:00