mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-18 04:46:06 +00:00
Code Issues Packages Projects Releases Wiki Activity
15,475 Commits 503 Branches 263 Tags
Commit Graph

10 Commits

Author SHA1 Message Date
Peter Hawkins
c2d6fcc0e6 Split core.py and several files in an SCC with it into a separate Bazel build target. 
PiperOrigin-RevId: 520192610
 2023-03-28 18:31:13 -07:00
Yash Katariya
a9e48af260 Deprecated xla_call_p since it has been replaced with pjit.pjit_p 
PiperOrigin-RevId: 518921538
 2023-03-23 11:44:42 -07:00
Yash Katariya
c2d5527f72 [Jax cleanup] 
* Remove lower_xla_callable and all related functions
* Remove pxla.device_put
* Remove dispatch.device_put_handlers

PiperOrigin-RevId: 517249345
 2023-03-16 15:47:28 -07:00
Peter Hawkins
8fb1fd318d Replace jax._src.util.prod with math.prod. 
math.prod() was added in Python 3.8, so we can assume it is always present.

PiperOrigin-RevId: 513011144
 2023-02-28 12:41:00 -08:00
Peter Hawkins
f66f6ec98a [JAX] Move jax._src.lib.xla_bridge to jax._src.xla_bridge. 
Limit jax._src.lib to shims around jaxlib and nothing else.

The goal of this change is to avoid a dependency cycle between the rest of jax and jax._src.lib in a Bazel build. This allows the types for jax._src.lib to be inferred by pytype in isolation without referring to the rest of JAX.

PiperOrigin-RevId: 512922397
 2023-02-28 07:01:57 -08:00
Peter Hawkins
148774587a Remove circular dependency between source_info_util and util. 
Move util.new_name_stack into source_info_util. Replace uses of util.extend_name_stack with stack.extend().

PiperOrigin-RevId: 512685810
 2023-02-27 11:41:46 -08:00
Peter Hawkins
54269c1145 Remove more exported names from jax.interpreters.xla. 
None of these appear to have public users, and this module is not included in the deprecation policy.

Also:
* shorten a number of alias chains.
* move make_op_metadata() into its only caller in jax2tf
* delete the unused function dtype_to_primitive_type.
PiperOrigin-RevId: 510205315
 2023-02-16 11:56:30 -08:00
Peter Hawkins
c6a99b699e Remove jax.interpreters.xla.lower_fun. 
This function has been a stub that does nothing useful for a long time, and the only user I can find is Equinox which already guards this with a hasattr(xla, 'lower_fun') guard.

PiperOrigin-RevId: 510142446
 2023-02-16 07:51:15 -08:00
Matthew Johnson
a964dc3b9a simpler pretty-print for pjit, tweak custom pp rule signature 2023-02-09 12:45:51 -08:00
Peter Hawkins
6860cb8d2a Move jax.interpreters.xla to jax._src.interpreters.xla. 
Replace jax.interpreters.xla with a shim that re-exports names that are likely to be used externally.

PiperOrigin-RevId: 507895040
 2023-02-07 15:01:32 -08:00