Michael Hudgins d4d1518c3d Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax ...

PiperOrigin-RevId: 676843138

2024-09-20 07:52:33 -07:00

..

control_flow

Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax

2024-09-20 07:52:33 -07:00

__init__.py

Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.".

2022-09-22 12:27:19 -07:00

ann.py

Clean up and fix primal type to tangent type mapping

2024-09-18 13:43:54 -07:00

convolution.py

Update the docs for conv_general_dilated to clarify 'W' 'H'.

2024-09-10 14:02:06 -07:00

eigh.py

Partial rollback of https://github.com/google/jax/pull/23353 as discussed in https://github.com/google/jax/pull/23353#issuecomment-2326604708

2024-09-03 09:49:22 -07:00

fft.py

Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax

2024-09-20 07:52:33 -07:00

lax.py

Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax

2024-09-20 07:52:33 -07:00

linalg.py

Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax

2024-09-20 07:52:33 -07:00

other.py

Update the docs for conv_general_dilated to clarify 'W' 'H'.

2024-09-10 14:02:06 -07:00

parallel.py

refactor jax.lax to not depend on jax.numpy

2024-09-01 07:49:49 -07:00

qdwh.py

Partial rollback of https://github.com/google/jax/pull/23353 as discussed in https://github.com/google/jax/pull/23353#issuecomment-2326604708

2024-09-03 09:49:22 -07:00

slicing.py

Clean up and fix primal type to tangent type mapping

2024-09-18 13:43:54 -07:00

special.py

Use MLIR generated convenience functions athing(...) instead of writing AThingOp(...).result.

2023-11-17 11:47:14 -08:00

stack.py

Partial rollback of https://github.com/google/jax/pull/23353 as discussed in https://github.com/google/jax/pull/23353#issuecomment-2326604708

2024-09-03 09:49:22 -07:00

svd.py

Partial rollback of https://github.com/google/jax/pull/23353 as discussed in https://github.com/google/jax/pull/23353#issuecomment-2326604708

2024-09-03 09:49:22 -07:00

utils.py

Initial integration of sharding in types in JAX. Currently we just support nary ops in forward only sharding propagation. Currently this functionality is experimental and hidden behind jax_sharding_in_types config flag.

2024-08-29 10:50:04 -07:00

windowed_reductions.py

Clean up and fix primal type to tangent type mapping

2024-09-18 13:43:54 -07:00