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

..

__init__.py

[sparse] implement bcsr_concatenate

2023-02-15 14:10:56 -08:00

_base.py

Replace jax._src.util.prod with math.prod.

2023-02-28 12:41:00 -08:00

_lowerings.py

[sparse] add low-level primitives wrapping cuda csr spmv and spmm.

2023-03-06 11:34:30 -08:00

ad.py

migrate more internal dependencies from jax.core to jax._src.core

2023-02-14 23:01:11 -08:00

api.py

migrate more internal dependencies from jax.core to jax._src.core

2023-02-14 23:01:11 -08:00

bcoo.py

Silence some pytype errors.

2023-03-09 16:40:56 -08:00

bcsr.py

[sparse] modify impl rules to always use lowering path

2023-03-02 14:27:18 -08:00

coo.py

Remove leading underscores in jax._src.numpy.util

2023-03-13 12:18:36 -07:00

csr.py

Remove leading underscores in jax._src.numpy.util

2023-03-13 12:18:36 -07:00

linalg.py

Replace uses of jnp.ndarray with jax.Array inside JAX.

2023-02-15 14:53:00 -08:00

random.py

[sparse] more careful handling of index_dtype in rand_bcoo

2023-03-07 08:34:10 -08:00

test_util.py

[sparse] add BCSR.to_bcoo and from_bcoo methods

2023-01-30 10:42:05 -08:00

transform.py

Deprecated xla_call_p since it has been replaced with pjit.pjit_p

2023-03-23 11:44:42 -07:00

util.py

migrate more internal dependencies from jax.core to jax._src.core

2023-02-14 23:01:11 -08:00