Comments (2)
hawkinsp
commented on June 1, 2024
Yes. That's correct and how it works at the moment. The custom partitioning ultimately refers to a Python object, which is why it's not stable run to run.
from jax.
nouiz
commented on June 1, 2024
If this only cover the python callback, could we just remove it from the key?
Not everything is 100% versioned right now (like XLA isn't versioned).
So it would just end up to the end user responsibility to make sure handle the cache correctly?
Or could we get the python ast of the callbacks and hash it? I see that as more work and not sure it is useful.
from jax.
Related Issues (20)
- RuntimeError: Bad StatusOr access: INTERNAL: Mosaic failed to compile TPU kernel: Bad lhs type in tpu.matmul HOT 7
- Unexpected scaling of sharding/pmap HOT 1
- Low calculation performance compared to autograd elementwise_grad HOT 1
- D2H (gpu -> cpu) transfer via `device_put` is very slow HOT 7
- example path HOT 1
- Tracking issue: inline docstrings HOT 2
- Slow initialisation on H100 HOT 1
- Unexpected Assertion error when using lax.with_sharding_constraint HOT 4
- Can´t get jax profiling to work HOT 2
- Cryptic error message when running pallas kernel on TPU HOT 3
- jax/jaxlib not working HOT 5
- `softmax` returns `NaN` when inputs are `inf` HOT 2
- jax.scipy.special.hyp1f1 unstable where scipy.special.hyp1f1 is not HOT 1
- pure_callback is not passing `np.ndarray` typed inputs HOT 7
- jax raise AttributeError on import when running with python doctring strip optimization
- jax-metal: gather fails with unsigned indices
- jax-metal: segmentation fault inside `jax.lax.while_loop` HOT 1
- jax-metal: 3-D convolution crashes
- constants in Pallas kernel causes `ValueError: safe_zip() argument 2 is shorter than argument 1`
- Involuntary full rematerialization for advanced data movements HOT 3
Recommend Projects
-
React
A declarative, efficient, and flexible JavaScript library for building user interfaces.
-
Vue.js
🖖 Vue.js is a progressive, incrementally-adoptable JavaScript framework for building UI on the web.
-
Typescript
TypeScript is a superset of JavaScript that compiles to clean JavaScript output.
-
TensorFlow
An Open Source Machine Learning Framework for Everyone
-
Django
The Web framework for perfectionists with deadlines.
-
Laravel
A PHP framework for web artisans
-
D3
Bring data to life with SVG, Canvas and HTML. 📊📈🎉
-
Recommend Topics
-
javascript
JavaScript (JS) is a lightweight interpreted programming language with first-class functions.
-
web
Some thing interesting about web. New door for the world.
-
server
A server is a program made to process requests and deliver data to clients.
-
Machine learning
Machine learning is a way of modeling and interpreting data that allows a piece of software to respond intelligently.
-
Visualization
Some thing interesting about visualization, use data art
-
Game
Some thing interesting about game, make everyone happy.
Recommend Org
-
Facebook
We are working to build community through open source technology. NB: members must have two-factor auth.
-
Microsoft
Open source projects and samples from Microsoft.
-
Google
Google ❤️ Open Source for everyone.
-
Alibaba
Alibaba Open Source for everyone
-
D3
Data-Driven Documents codes.
-
Tencent
China tencent open source team.
from jax.