how to use pytensor with jitted functions in jax (and tracer arrays)

[kratsg]

I'm working on a project that builds likelihood functions, using pytensor to compose these expressions and then pass in the values straight into the compiled function. For the following (somewhat obscure code, but hopefully enlightening enough to understand what I'm trying to do)

import pyhs3
import jax
import jax.numpy as jnp


ws = pyhs3.Workspace.load('tests/test_pdf/rf501_simultaneouspdf.json')
model = ws.model(mode='jax')

@jax.jit
def nll(pars):
  return -2 * model.logpdf('model', x=jnp.array(5.0), **pars)

# all your parameters...
pars = jax.tree.map(jnp.asarray, {'f': 0.2, 'mean': 0.0, 'sigma': 0.3, 'mean2': 0.0, 'sigma2': 0.3})
print("NLL", nll(pars))

will fail with

Traceback (most recent call last):
  File "/Users/kratsg/pyhs3/optimistix_test.py", line 35, in <module>
    print("NLL", nll(pars))
                 ~~~^^^^^^
  File "/Users/kratsg/pyhs3/optimistix_test.py", line 31, in nll
    return -2 * model.logpdf('model', x=jnp.array(5.0), **pars)
                ~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/kratsg/pyhs3/src/pyhs3/core.py", line 572, in logpdf
    return np.log(self.pdf(name, **parametervalues))
                  ~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/kratsg/pyhs3/src/pyhs3/core.py", line 522, in pdf
    return func(*positional_values)
  File "/Users/kratsg/.pyenv/versions/pyhs3/lib/python3.13/site-packages/pytensor/compile/function/types.py", line 979, in __call__
    raise
  File "/Users/kratsg/.pyenv/versions/pyhs3/lib/python3.13/site-packages/pytensor/tensor/type.py", line 228, in filter
    converted_data = np.asarray(data, self.dtype)
jax.errors.TracerArrayConversionError: Bad input argument to pytensor function with name "/Users/kratsg/pyhs3/src/pyhs3/core.py:489" at index 0 (0-based).
Backtrace when that variable is created:

  File "/Users/kratsg/pyhs3/optimistix_test.py", line 27, in <module>
    model = ws.model(mode='jax')
  File "/Users/kratsg/pyhs3/src/pyhs3/core.py", line 249, in model
    return Model(
  File "/Users/kratsg/pyhs3/src/pyhs3/core.py", line 323, in __init__
    self._build_dependency_graph(functions, distributions, progress)
  File "/Users/kratsg/pyhs3/src/pyhs3/core.py", line 407, in _build_dependency_graph
    self.parameters[node_name] = create_bounded_tensor(
  File "/Users/kratsg/pyhs3/src/pyhs3/core.py", line 708, in create_bounded_tensor
    tensor = kind(name)
The numpy.ndarray conversion method __array__() was called on traced array with shape float32[]
The error occurred while tracing the function nll at /Users/kratsg/pyhs3/optimistix_test.py:29 for jit. This concrete value was not available in Python because it depends on the value of the argument pars['f'].
See https://docs.jax.dev/en/latest/errors.html#jax.errors.TracerArrayConversionError
--------------------
For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.

because in

pytensor/pytensor/tensor/type.py

Line 153 in 1d825dd

def filter(self, data, strict=False, allow_downcast=None) -> np.ndarray:

(and really everywhere in this file) are np.array / np.asarray calls this makes things not work well with jax.

model.distributions['model']
# Add.0
model._compiled_functions['model']
# <pytensor.compile.function.types.Function object at 0x110913380>
pp(model.distributions['model'])
# '((0.0 + (Clip(f, 0.0, 1.0) * ((1.0 / (Sqrt(6.283185307179586) * Clip(sigma, 0.1, 10.0))) * Exp((-0.5 * (((Clip(x, -8.0, 8.0) - Clip(mean, -8.0, 8.0)) / Clip(sigma, 0.1, 10.0)) ** 2)))))) + ((1 - (0.0 + Clip(f, 0.0, 1.0))) * ((1.0 / (Sqrt(6.283185307179586) * Clip(sigma2, 0.1, 10.0))) * Exp((-0.5 * (((Clip(x, -8.0, 8.0) - Clip(mean2, -3.0, 3.0)) / Clip(sigma2, 0.1, 10.0)) ** 2))))))'

where the function is compiled like

            self._compiled_functions[name] = cast(
                Callable[..., npt.NDArray[np.float64]],
                function(
                    inputs=inputs,
                    outputs=combined_expression,
                    mode=compilation_mode,
                    on_unused_input="ignore",
                ),
            )

so I guess it's not clear how to compile something for jax so that it can be used in a jit'd context like the above code.

./cc @pfackeldey

[ricardoV94]

In this case you don't want to compile a PyTensor function, but transpile the underlying graph to JAX. PyMC does this when working with JAX based simplers here: https://github.com/pymc-devs/pymc/blob/340e403b8813ab5f3699a476cc828cc92c4f9d50/pymc/sampling/jax.py#L122-L133

If you have shared variables you have to make a decision how to handle them. In PyMC we treat them as constants as we know users wouldn't be updating them during the time the JAX samplers are working. But that is context-specific. If you don't have shared variables then it's just those 3 lines really:

def get_jaxified_graph(
    inputs: list[TensorVariable] | None = None,
    outputs: list[TensorVariable] | None = None,
) -> Callable[[list[TensorVariable]], list[TensorVariable]]:
    """Compile a PyTensor graph into an optimized JAX function."""
    graph = _replace_shared_variables(outputs) if outputs is not None else None

    fgraph = FunctionGraph(inputs=inputs, outputs=graph, clone=True)
    mode.JAX.optimizer.rewrite(fgraph)

    # We now jaxify the optimized fgraph
    return jax_funcify(fgraph)

If you want to define the graph dynamically when jax calls JIT you'll have to do some extra work to convert the JAX traced arrays to pytensor variables, define the pytensor graph and then return the jax graph. Couldn't tell whether you were in that context from your description

[ricardoV94]

A pytensor.function compiled to JAX backend is not a jax function that JAX will be happy to trace through.

Specifically there's a way we handle inputs and outputs that is a bit archaic. We store values in shared lists and then pop them out for the jitted function, and then retrieve the outputs and put them in other lists. We also do some validation to ensure the inputs/outputs are numpy arrays, as users using a pytensor function are supposed to work with those, regardless of the backend.

You however want to skip this middle layer, and work directly with JAX which is what the code I shared above does. It will return a regular python function that is comprised exclusively of JAX operations. It will happily accept jax arrays and output jax arrays.

Shared variables are a specific pytensor concept, usually used to manage some state across function calls: https://pytensor.readthedocs.io/en/latest/tutorial/examples.html#using-shared-variables. This is something extra that pytensor.function can handle for the user, but which makes no sense to the underlying jax function. Hence why I brought it up. If you are not creating any pytensor.shared, or using random variables in the graph you don't have to worry about those.

Hope that clarifies things a bit.

[kratsg]

Here's a working gist :) https://gist.github.com/kratsg/2c83d4dd3367664554ba842e85c48879 -- thanks for the pointers. I'm not sure I would have dug this out of pytensor's extensive documentation. I changed the function around so for me, it looks like

def get_jaxified_graph(
    model: pyhs3.Model,
    name: str,
) -> tuple[list[TensorVar], Callable[[list[TensorVar]], list[TensorVar]]]:
    """Compile a PyTensor graph into an optimized JAX function."""

    dist = model.distributions[name]
    inputs: list[TensorVar] = [
        var for var in graph_inputs([dist]) if var.name is not None
    ]

    fgraph = FunctionGraph(inputs=inputs, outputs=[dist], clone=True)
    mode.JAX.optimizer.rewrite(fgraph)

    # We now jaxify the optimized fgraph
    return inputs, jax_funcify(fgraph)

but this works quite well :)

[ricardoV94]

looks good. for inputs you probably want explicit_graph_inputs (to filter out constants). In theory a constant could have a name as well.

[ricardoV94]

But you actually don't need to specify the inputs if you do FunctionGraph(outputs=[dist]) it will find the inputs for you. One thing I would be concerned is that the inputs are in the right order when called from jax later

[kratsg]

I'll make a note of explicit_graph_inputs and cross-reference this discussion so I can come bak to it later to improve the API on our end.

But you actually don't need to specify the inputs if you do FunctionGraph(outputs=[dist]) it will find the inputs for you.

That's nifty, I guess because it's doing effectively the same logic that I was doing, just more baked in (so that really does simplify some of the work).

One thing I would be concerned is that the inputs are in the right order when called from jax later

yes, I added some code to properly stitch the order of the inputs, which is at least why I am still figuring out the best way to pass that information around (both the function and it's signature).