Rewrite size input of RandomVariables in JAX backend

Author: rlouf

pytensor/compile/mode.py

@@ -449,7 +449,7 @@ def clone(self, link_kwargs=None, optimizer="", **kwargs):
 
 JAX = Mode(
  JAXLinker(),
- RewriteDatabaseQuery(include=["fast_run"], exclude=["cxx_only", "BlasOpt"]),
+ RewriteDatabaseQuery(include=["fast_run", "jax"], exclude=["cxx_only", "BlasOpt"]),
 )
 NUMBA = Mode(
  NumbaLinker(),

pytensor/link/jax/dispatch/random.py

@@ -8,6 +8,7 @@
 
 import pytensor.tensor.random.basic as aer
 from pytensor.link.jax.dispatch.basic import jax_funcify, jax_typify
+from pytensor.link.jax.dispatch.shape import JAXShapeTuple
 from pytensor.tensor.shape import Shape, Shape_i
 
 
@@ -28,7 +29,7 @@
 
 
 def assert_size_argument_jax_compatible(node):
- """Assert whether the current node can be compiled.
+ """Assert whether the current node can be JIT-compiled by JAX.
 
  JAX can JIT-compile `jax.random` functions when the `size` argument
  is a concrete value, i.e. either a constant or the shape of any
@@ -37,7 +38,7 @@ def assert_size_argument_jax_compatible(node):
  """
  size = node.inputs[1]
  size_op = size.owner.op
- if not isinstance(size_op, (Shape, Shape_i)):
+ if not isinstance(size_op, (Shape, Shape_i, JAXShapeTuple)):
  raise NotImplementedError(SIZE_NOT_COMPATIBLE)
 
 

pytensor/link/jax/dispatch/shape.py

@@ -1,8 +1,31 @@
 import jax.numpy as jnp
 
 from pytensor.graph import Constant
+from pytensor.graph.basic import Apply
+from pytensor.graph.op import Op
 from pytensor.link.jax.dispatch.basic import jax_funcify
 from pytensor.tensor.shape import Reshape, Shape, Shape_i, SpecifyShape, Unbroadcast
+from pytensor.tensor.type import TensorType
+
+
+class JAXShapeTuple(Op):
+ """Dummy Op that represents a `size` specified as a tuple."""
+
+ def make_node(self, *inputs):
+ dtype = inputs[0].type.dtype
+ otype = TensorType(dtype, shape=(len(inputs),))
+ return Apply(self, inputs, [otype()])
+
+ def perform(self, *inputs):
+ return tuple(inputs)
+
+
+@jax_funcify.register(JAXShapeTuple)
+def jax_funcify_JAXShapeTuple(op, **kwargs):
+ def shape_tuple_fn(*x):
+ return tuple(x)
+
+ return shape_tuple_fn
 
 
 @jax_funcify.register(Reshape)

pytensor/tensor/random/rewriting/__init__.py

@@ -1,2 +1,10 @@
 # TODO: This is for backward-compatibility; remove when reasonable.
 from pytensor.tensor.random.rewriting.basic import *
+
+
+# isort: off
+
+# Register JAX specializations
+import pytensor.tensor.random.rewriting.jax
+
+# isort: on

pytensor/tensor/random/rewriting/jax.py

@@ -0,0 +1,52 @@
+from pytensor.compile import optdb
+from pytensor.graph.rewriting.basic import in2out, node_rewriter
+from pytensor.tensor.basic import MakeVector
+from pytensor.tensor.elemwise import DimShuffle
+from pytensor.tensor.random.op import RandomVariable
+
+
+@node_rewriter([RandomVariable])
+def size_parameter_as_tuple(fgraph, node):
+ """Replace `MakeVector` and `DimShuffle` (when used to transform a scalar
+ into a 1d vector) when they are found as the input of a `size` or `shape`
+ parameter by `JAXShapeTuple` during transpilation.
+
+ The JAX implementations of `MakeVector` and `DimShuffle` always return JAX
+ `TracedArrays`, but JAX only accepts concrete values as inputs for the `size`
+ or `shape` parameter. When these `Op`s are used to convert scalar or tuple
+ inputs, however, we can avoid tracing by making them return a tuple of their
+ inputs instead.
+
+ Note that JAX does not accept scalar inputs for the `size` or `shape`
+ parameters, and this rewrite also ensures that scalar inputs are turned into
+ tuples during transpilation.
+
+ """
+ from pytensor.link.jax.dispatch.shape import JAXShapeTuple
+
+ size_arg = node.inputs[1]
+ size_node = size_arg.owner
+
+ if size_node is None:
+ return
+
+ if isinstance(size_node.op, JAXShapeTuple):
+ return
+
+ if isinstance(size_node.op, MakeVector) or (
+ isinstance(size_node.op, DimShuffle)
+ and size_node.op.input_broadcastable == ()
+ and size_node.op.new_order == ("x",)
+ ):
+ # Here PyTensor converted a tuple or list to a tensor
+ new_size_args = JAXShapeTuple()(*size_node.inputs)
+ new_inputs = list(node.inputs)
+ new_inputs[1] = new_size_args
+
+ new_node = node.clone_with_new_inputs(new_inputs)
+ return new_node.outputs
+
+
+optdb.register(
+ "jax_size_parameter_as_tuple", in2out(size_parameter_as_tuple), "jax", position=100
+)

tests/link/jax/test_basic.py

@@ -27,7 +27,7 @@ def set_pytensor_flags():
 jax = pytest.importorskip("jax")
 
 
-opts = RewriteDatabaseQuery(include=[None], exclude=["cxx_only", "BlasOpt"])
+opts = RewriteDatabaseQuery(include=["jax"], exclude=["cxx_only", "BlasOpt"])
 jax_mode = Mode(JAXLinker(), opts)
 py_mode = Mode("py", opts)
 

tests/link/jax/test_random.py

@@ -454,25 +454,44 @@ def test_random_concrete_shape():
  assert jax_fn(np.ones((2, 3))).shape == (2, 3)
 
 
-@pytest.mark.xfail(reason="size argument specified as a tuple is a `DimShuffle` node")
 def test_random_concrete_shape_subtensor():
+ """JAX should compile when a concrete value is passed for the `size` parameter.
+
+ This test ensures that the `DimShuffle` `Op` used by PyTensor to turn scalar
+ inputs into 1d vectors is replaced by an `Op` that turns concrete scalar
+ inputs into tuples of concrete values using the `jax_size_parameter_as_tuple`
+ rewrite.
+
+ JAX does not accept scalars as `size` or `shape` arguments, so this is a
+ slight improvement over their API.
+
+ """
  rng = shared(np.random.RandomState(123))
  x_at = at.dmatrix()
  out = at.random.normal(0, 1, size=x_at.shape[1], rng=rng)
  jax_fn = function([x_at], out, mode=jax_mode)
  assert jax_fn(np.ones((2, 3))).shape == (3,)
 
 
-@pytest.mark.xfail(reason="size argument specified as a tuple is a `MakeVector` node")
 def test_random_concrete_shape_subtensor_tuple():
+ """JAX should compile when a tuple of concrete values is passed for the `size` parameter.
+
+ This test ensures that the `MakeVector` `Op` used by PyTensor to turn tuple
+ inputs into 1d vectors is replaced by an `Op` that turns a tuple of concrete
+ scalar inputs into tuples of concrete values using the
+ `jax_size_parameter_as_tuple` rewrite.
+
+ """
  rng = shared(np.random.RandomState(123))
  x_at = at.dmatrix()
  out = at.random.normal(0, 1, size=(x_at.shape[0],), rng=rng)
  jax_fn = function([x_at], out, mode=jax_mode)
  assert jax_fn(np.ones((2, 3))).shape == (2,)
 
 
-@pytest.mark.xfail(reason="`size_at` should be specified as a static argument")
+@pytest.mark.xfail(
+ reason="`size_at` should be specified as a static argument", strict=True
+)
 def test_random_concrete_shape_graph_input():
  rng = shared(np.random.RandomState(123))
  size_at = at.scalar()