pyhs3.JaxifiedGraph

class pyhs3.JaxifiedGraph(inputs, input_names, fn)

A JAX-callable wrapper around a compiled PyTensor expression.

Produced by jaxify(). Supports keyword-argument calls (the primary interface for dict-pytree-based optimizers like optimistix or everwillow) and positional calls via call_positional().

inputs

Tuple of PyTensor input variables, in evaluation order.

input_names

Names of those variables (same order as inputs).

fn

The raw JAX callable returned by jax_funcify.

Parameters:

• inputs (tuple[TypeAliasType, ...])

• input_names (tuple[str, ...])

• fn (Callable[..., tuple])

__init__(inputs, input_names, fn)

Parameters:

• inputs (tuple[TypeAliasType, ...])

• input_names (tuple[str, ...])

• fn (Callable[..., tuple])

Methods

__init__(inputs, input_names, fn)
call_positional(*args)	Call with positional arguments — pure passthrough to fn.

Attributes

inputs
input_names
fn

JaxifiedGraph.call_positional(*args)

Call with positional arguments — pure passthrough to fn.

Parameters:

*args (object) – Values in the same order as self.input_names.

Return type:

Whatever fn returns (typically a 1-tuple of JAX arrays).

JaxifiedGraph.inputs: tuple[TensorVar, ...]

JaxifiedGraph.input_names: tuple[str, ...]

JaxifiedGraph.fn: Callable[[...], tuple]