Rewrite size input of RandomVariables in JAX backend

Author: rlouf

aesara/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(),

aesara/link/jax/dispatch/__init__.py

@@ -13,4 +13,7 @@
 import aesara.link.jax.dispatch.elemwise
 import aesara.link.jax.dispatch.scan
 
+# Load specialized rewrites
+import aesara.link.jax.dispatch.rewriting
+
 # isort: on

aesara/link/jax/dispatch/random.py

@@ -8,6 +8,7 @@
 
 import aesara.tensor.random.basic as aer
 from aesara.link.jax.dispatch.basic import jax_funcify, jax_typify
+from aesara.link.jax.dispatch.shape import JAXShapeTuple
 from aesara.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)
 
 

aesara/link/jax/dispatch/rewriting.py

@@ -0,0 +1,35 @@
+from aesara.graph.rewriting.basic import node_rewriter, in2out
+from aesara.compile import optdb
+from aesara.link.jax.dispatch.shape import JAXShapeTuple
+from aesara.tensor.basic import MakeVector
+from aesara.tensor.elemwise import DimShuffle
+from aesara.tensor.random.op import RandomVariable
+
+
+@node_rewriter([RandomVariable])
+def size_parameter_as_tuple(fgraph, node):
+    if isinstance(node.op, RandomVariable):
+        size_arg = node.inputs[1]
+        if isinstance(size_arg.owner.op, MakeVector):
+            # Here Aesara converted a tuple or list to a tensor
+            new_size_args = JAXShapeTuple()(*size_arg.owner.inputs)
+            node.inputs[1] = new_size_args
+            return node.outputs
+
+        if isinstance(size_arg.owner.op, DimShuffle):
+            # Here Aesara is turning a scalar input into a 1d vector
+            if (
+                size_arg.owner.op.input_broadcastable == ()
+                and size_arg.owner.op.new_order == ("x",)
+            ):
+                new_size_args = JAXShapeTuple()(*size_arg.owner.inputs)
+                node.inputs[1] = new_size_args
+                return node.outputs
+
+
+optdb.register(
+    "jax_size_parameter_as_tuple",
+    in2out(size_parameter_as_tuple),
+    "jax",
+    position=100
+)

aesara/link/jax/dispatch/shape.py

@@ -1,8 +1,34 @@
 import jax.numpy as jnp
 
 from aesara.graph import Constant
+from aesara.graph.basic import Apply
+from aesara.graph.op import Op
 from aesara.link.jax.dispatch.basic import jax_funcify
 from aesara.tensor.shape import Reshape, Shape, Shape_i, SpecifyShape, Unbroadcast
+from aesara.tensor.type import TensorType
+
+
+class JAXShapeTuple(Op):
+    """Dummy Op that represents a `size` specified as a tuple."""
+
+    def __init__(self):
+        pass
+
+    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)

tests/link/jax/test_random.py

@@ -454,7 +454,7 @@ 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")
+# @pytest.mark.xfail(reason="size argument specified as a tuple is a `DimShuffle` node")
 def test_random_concrete_shape_subtensor():
     rng = shared(np.random.RandomState(123))
     x_at = at.dmatrix()
@@ -463,7 +463,7 @@ def test_random_concrete_shape_subtensor():
     assert jax_fn(np.ones((2, 3))).shape == (3,)
 
 
-@pytest.mark.xfail(reason="size argument specified as a tuple is a `MakeVector` node")
+# @pytest.mark.xfail(reason="size argument specified as a tuple is a `MakeVector` node")
 def test_random_concrete_shape_subtensor_tuple():
     rng = shared(np.random.RandomState(123))
     x_at = at.dmatrix()