Comments (2)
@HarshvirSandhu thanks for reporting the issue.
That tutorial was written back in a time when JAX + jitting was more flexible. These days all jitted functions must have constant shape, which means a graph like the one in the example can never be translated to JAX, since it's fundamentally a function with dynamic shapes. The first error can be avoided by specifying the scalar has an integer dtype, but that will only kick the can further. After also updating the dispatch function signature we get this:
import jax.numpy as jnp
from pytensor.tensor.basic import Eye
from pytensor.link.jax.dispatch import jax_funcify
from tests.link.jax.test_basic import compare_jax_and_py
from pytensor.graph import FunctionGraph
import pytensor.tensor as pt
@jax_funcify.register(Eye)
def jax_funcify_Eye(op, node, **kwargs):
dtype = op.dtype
def eye(N, M, k):
return jnp.eye(N, M, k, dtype=dtype)
return eye
def test_jax_Eye():
"""Test JAX conversion of the `Eye` `Op`."""
x_at = pt.scalar(dtype="int")
eye_var = pt.eye(x_at)
out_fg = FunctionGraph(outputs=[eye_var])
compare_jax_and_py(out_fg, [3])
test_jax_Eye()
TypeError: Shapes must be 1D sequences of concrete values of integer type, got (Traced<ShapedArray(int64[])>with<DynamicJaxprTrace(level=1/0)>,). 'N' argument of jnp.eye().
If using `jit`, try using `static_argnums` or applying `jit` to smaller subfunctions.
We should use a different example and perhaps keep this in a section of its own explaining that not all PyTensor graphs can be converted into a jitted JAX-graph, usually those that have dynamic shapes.
from pytensor.
Hi,
I would like to take up this issue as part of my GSoC application. I will open a PR on this after changing to a more suitable example.
from pytensor.
Related Issues (20)
- Upload codecov job is failing systematically
- Reconcile environment for development and docs building
- Handling tensor with variable length of dimension
- Add rewrite to merge multiple `SVD` `Op`s with different settings HOT 4
- Codecov is failing to update due to limit without token
- Can we automate transfer of stack traces?
- A graph replace that changes the type from dynamic to static can lead to miscompilation HOT 2
- Fusion rewriter error HOT 1
- ENH: Implement `pt.pad`
- Index page of the documentation has the wrong header link to github
- Get rid of join view flag
- Flaky `tests/tensor/rewriting/test_linalg.py::test_local_lift_through_linalg[kron-pinv-batched]` HOT 1
- Add option to implement vectorized hessian/jacobian
- BUG: `pt.full` doesn't work with int shape HOT 3
- ENH: Detect clang++ as a g++ alternative on Linux HOT 2
- Don't run mypy in pre-commit HOT 2
- BUG: <```uint64 dtype``` is broken for ```Max```> HOT 3
- Handle more AdvancedSetSubtensor ops in numba HOT 1
- Update xfail reason "Omnistaging can't be disabled" in JAX tests
- Missed scan rewrites
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 pytensor.