Jake VanderPlas 39dea4c29e [sparse] modify impl rules to always use lowering path

2023-03-02 14:27:18 -08: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 SpMV & SpMM

2023-03-01 05:56:31 -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

[sparse] modify impl rules to always use lowering path

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

bcsr.py

[sparse] modify impl rules to always use lowering path

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

coo.py

[sparse] modify impl rules to always use lowering path

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

csr.py

[sparse] modify impl rules to always use lowering path

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

linalg.py

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

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

random.py

Switch to new-style f-strings

2022-12-01 09:14:16 -08:00

test_util.py

[sparse] add BCSR.to_bcoo and from_bcoo methods

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

transform.py

[sparse] implement bcsr_concatenate

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

util.py

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

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