Search results

Suggest tfp-nighly

jax

no releasenotes

Status: Open (in progress).

pymc-devs/pytensornumber 1763

#1763 In pymc-devs/pytensor;

· ricardoV94 opened on Dec 1, 2025

Fix <code>wrap_jax</code> when there is a mix of statically known and unknown shapes

bugSomething isn't workingSomething isn't working

jax

Status: Open (in progress).

pymc-devs/pytensornumber 1722

#1722 In pymc-devs/pytensor;

· ricardoV94 opened on Nov 16, 2025

Create <code>vectorized</code>, <code>value_and_grad</code> and <code>shape</code> versions of <code>JAXOp</code>

enhancementNew feature or requestNew feature or request

jax

Op implementation

Status: Open.

#1645 In pymc-devs/pytensor;

· ricardoV94 opened on Oct 9, 2025

Run jax doctests in jax CI

docs

GitHub CI/CD

jax

Status: Open.

#1624 In pymc-devs/pytensor;

· ricardoV94 opened on Sep 29, 2025

Add error message to <code>NonZero</code> Op telling users it cannot be jitted in JAX

backend compatibility

beginner friendly

help wantedExtra attention is neededExtra attention is needed

jax

Status: Open.

#1509 In pymc-devs/pytensor;

· jessegrabowski opened on Jul 1, 2025

Transform graph to make large constants symbolic inputs

enhancementNew feature or requestNew feature or request

jax

Status: Open.

#1224 In pymc-devs/pytensor;

· ricardoV94 opened on Feb 19, 2025

Implement <code>@as_jax_op</code> to wrap a JAX function for use in PyTensor

enhancementNew feature or requestNew feature or request

jax

Op implementation

Status: Draft (not ready).

pymc-devs/pytensornumber 1120

#1120 In pymc-devs/pytensor;

· jdehning opened on Dec 12, 2024

Two parametrization of JAX benchmark test are taking hours

jax

needs infoAdditional information requiredAdditional information required

tests

Status: Open.

#1060 In pymc-devs/pytensor;

· ricardoV94 opened on Oct 31, 2024

Use <code>expand_dims</code> / <code>squeeze</code> in JAX implementation of <code>Dimshuffle</code>

beginner friendly

jax

Status: Open.

#847 In pymc-devs/pytensor;

· ricardoV94 opened on Jun 24, 2024

minor bug: jax tests reference jax.interpreters.xla.DeviceArray

beginner friendly

bugSomething isn't workingSomething isn't working

jax

tests

Status: Open.

#660 In pymc-devs/pytensor;

· leventov opened on Mar 7, 2024

Implement helper <code>@as_jax_op</code> to wrap JAX functions in PyTensor

backend compatibility

feature request

jax

torchPyTorch backendPyTorch backend

Status: Open.

#537 In pymc-devs/pytensor;

· ricardoV94 opened on Dec 7, 2023

No JAX dispatch for <code>mul_without_zeros</code>

backend compatibility

enhancementNew feature or requestNew feature or request

jax

Status: Open.

#526 In pymc-devs/pytensor;

· jessegrabowski opened on Nov 30, 2023