From 65917929742d0ddc934066216df41cee3f37b8f0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?R=C3=A9mi=20Louf?= Date: Wed, 7 Dec 2022 12:19:38 +0100 Subject: [PATCH 01/14] Raise when an array is resized with a boolean mask --- aesara/link/jax/dispatch/subtensor.py | 25 +++++++++++++++++++++- aesara/link/jax/dispatch/test_subtensor.py | 24 ++++++++++++++------- 2 files changed, 40 insertions(+), 9 deletions(-) diff --git a/aesara/link/jax/dispatch/subtensor.py b/aesara/link/jax/dispatch/subtensor.py index 822d78a6fa..b0c9644c9c 100644 --- a/aesara/link/jax/dispatch/subtensor.py +++ b/aesara/link/jax/dispatch/subtensor.py @@ -13,10 +13,33 @@ from aesara.tensor.type_other import MakeSlice +BOOLEAN_MASK_ERROR = """JAX does not support resizing arrays with boolean +masks. In some cases, however, it is possible to re-express your model +in a form that JAX can compile: + +>>> import aesara.tensor as at +>>> x_at = at.vector('x') +>>> y_at = x_at[x_at > 0].sum() + +can be re-expressed as: + +>>> import aesara.tensor as at +>>> x_at = at.vector('x') +>>> y_at = at.where(x_at > 0, x_at, 0).sum() +""" + + +def assert_indices_jax_compatible(node): + ilist = node.inputs[1] + if ilist.type.dtype == "bool": + raise NotImplementedError(BOOLEAN_MASK_ERROR) + + @jax_funcify.register(Subtensor) @jax_funcify.register(AdvancedSubtensor) @jax_funcify.register(AdvancedSubtensor1) -def jax_funcify_Subtensor(op, **kwargs): +def jax_funcify_Subtensor(op, node, **kwargs): + assert_indices_jax_compatible(node) idx_list = getattr(op, "idx_list", None) diff --git a/aesara/link/jax/dispatch/test_subtensor.py b/aesara/link/jax/dispatch/test_subtensor.py index 22cc492402..94d1244f24 100644 --- a/aesara/link/jax/dispatch/test_subtensor.py +++ b/aesara/link/jax/dispatch/test_subtensor.py @@ -47,16 +47,24 @@ def test_jax_Subtensors(): compare_jax_and_py(out_fg, []) -@pytest.mark.xfail( - version_parse(jax.__version__) >= version_parse("0.2.12"), - reason="Omnistaging cannot be disabled", -) -def test_jax_Subtensors_omni(): - x_at = at.arange(3 * 4 * 5).reshape((3, 4, 5)) - - # Boolean indices +def test_jax_Subtensor_boolean_mask(): + """JAX does not support resizing arrays with boolean masks.""" + x_at = at.arange(-5, 5) out_at = x_at[x_at < 0] assert isinstance(out_at.owner.op, at_subtensor.AdvancedSubtensor) + + with pytest.raises(NotImplementedError): + out_fg = FunctionGraph([], [out_at]) + compare_jax_and_py(out_fg, []) + + +@pytest.mark.xfail( + reason="Re-expressible boolean logic. We need a rewrite Aesara-side." +) +def test_jax_Subtensor_boolean_mask_reexpressible(): + """Some boolean logic can be re-expressed and JIT-compiled""" + x_at = at.arange(-5, 5) + out_at = x_at[x_at < 0].sum() out_fg = FunctionGraph([], [out_at]) compare_jax_and_py(out_fg, []) From 72930c925f45b943ad7a8f3ce96dc706cd8e22e4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?R=C3=A9mi=20Louf?= Date: Wed, 7 Dec 2022 15:32:11 +0100 Subject: [PATCH 02/14] Use constants or raise in JAX `Arange` implementation --- aesara/link/jax/dispatch/tensor_basic.py | 33 +++++++++++++++++++++--- tests/link/jax/test_tensor_basic.py | 15 ++++++++--- 2 files changed, 41 insertions(+), 7 deletions(-) diff --git a/aesara/link/jax/dispatch/tensor_basic.py b/aesara/link/jax/dispatch/tensor_basic.py index c15233175f..a578229fd7 100644 --- a/aesara/link/jax/dispatch/tensor_basic.py +++ b/aesara/link/jax/dispatch/tensor_basic.py @@ -1,5 +1,6 @@ import jax.numpy as jnp +from aesara.graph.basic import Constant from aesara.link.jax.dispatch.basic import jax_funcify from aesara.tensor.basic import ( Alloc, @@ -15,6 +16,15 @@ ) +ARANGE_CONCRETE_VALUE_ERROR = """JAX requires the arguments of `jax.numpy.arange` +to be constants. The graph that you defined thus cannot be JIT-compiled +by JAX. An example of a graph that can be compiled to JAX: + +>>> import aesara.tensor basic +>>> at.arange(1, 10, 2) +""" + + @jax_funcify.register(AllocDiag) def jax_funcify_AllocDiag(op, **kwargs): offset = op.offset @@ -43,9 +53,26 @@ def alloc(x, *shape): @jax_funcify.register(ARange) -def jax_funcify_ARange(op, **kwargs): - # XXX: This currently requires concrete arguments. - def arange(start, stop, step): +def jax_funcify_ARange(op, node, **kwargs): + """Register a JAX implementation for `ARange`. + + `jax.numpy.arange` requires concrete values for its arguments. Here we check + that the arguments are constant, and raise otherwise. + + TODO: Handle other situations in which values are concrete (shape of an array). + + """ + arange_args = node.inputs + constant_args = [] + for arg in arange_args: + if not isinstance(arg, Constant): + raise NotImplementedError(ARANGE_CONCRETE_VALUE_ERROR) + + constant_args.append(arg.value) + + start, stop, step = constant_args + + def arange(*_): return jnp.arange(start, stop, step, dtype=op.dtype) return arange diff --git a/tests/link/jax/test_tensor_basic.py b/tests/link/jax/test_tensor_basic.py index ef9738e0cd..696991ebad 100644 --- a/tests/link/jax/test_tensor_basic.py +++ b/tests/link/jax/test_tensor_basic.py @@ -52,15 +52,22 @@ def test_jax_MakeVector(): compare_jax_and_py(x_fg, []) -@pytest.mark.xfail(reason="jax.numpy.arange requires concrete inputs") +def test_arange(): + out = at.arange(1, 10, 2) + fgraph = FunctionGraph([], [out]) + compare_jax_and_py(fgraph, []) + + def test_arange_nonconcrete(): + """JAX cannot JIT-compile `jax.numpy.arange` when arguments are not concrete values.""" a = scalar("a") a.tag.test_value = 10 - out = at.arange(a) - fgraph = FunctionGraph([a], [out]) - compare_jax_and_py(fgraph, [get_test_value(i) for i in fgraph.inputs]) + + with pytest.raises(NotImplementedError): + fgraph = FunctionGraph([a], [out]) + compare_jax_and_py(fgraph, [get_test_value(i) for i in fgraph.inputs]) def test_jax_Join(): From e09919879dc7ce39ec74e9cee198341e7654db1d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?R=C3=A9mi=20Louf?= Date: Wed, 7 Dec 2022 15:44:40 +0100 Subject: [PATCH 03/14] Raise when trying to slice with a dynamic length --- aesara/link/jax/dispatch/subtensor.py | 30 ++++++++++++++------- aesara/link/jax/dispatch/test_subtensor.py | 31 +++++++++++++++++++--- 2 files changed, 48 insertions(+), 13 deletions(-) diff --git a/aesara/link/jax/dispatch/subtensor.py b/aesara/link/jax/dispatch/subtensor.py index b0c9644c9c..46293e2300 100644 --- a/aesara/link/jax/dispatch/subtensor.py +++ b/aesara/link/jax/dispatch/subtensor.py @@ -28,31 +28,43 @@ >>> y_at = at.where(x_at > 0, x_at, 0).sum() """ +DYNAMIC_SLICE_LENGTH_ERROR = """JAX does not support slicing arrays with a dynamic +slice length. +""" + + +def assert_indices_jax_compatible(node, idx_list): + from aesara.graph.basic import Constant + from aesara.tensor.var import TensorVariable -def assert_indices_jax_compatible(node): - ilist = node.inputs[1] - if ilist.type.dtype == "bool": - raise NotImplementedError(BOOLEAN_MASK_ERROR) + ilist = indices_from_subtensor(node.inputs[1:], idx_list) + for idx in ilist: + + if isinstance(idx, TensorVariable): + if idx.type.dtype == "bool": + raise NotImplementedError(BOOLEAN_MASK_ERROR) + elif isinstance(idx, slice): + for slice_arg in (idx.start, idx.stop, idx.step): + if slice_arg is not None and not isinstance(slice_arg, Constant): + raise NotImplementedError(DYNAMIC_SLICE_LENGTH_ERROR) @jax_funcify.register(Subtensor) @jax_funcify.register(AdvancedSubtensor) @jax_funcify.register(AdvancedSubtensor1) def jax_funcify_Subtensor(op, node, **kwargs): - assert_indices_jax_compatible(node) idx_list = getattr(op, "idx_list", None) + assert_indices_jax_compatible(node, idx_list) - def subtensor(x, *ilists): - + def subtensor_constant(x, *ilists): indices = indices_from_subtensor(ilists, idx_list) - if len(indices) == 1: indices = indices[0] return x.__getitem__(indices) - return subtensor + return subtensor_constant @jax_funcify.register(IncSubtensor) diff --git a/aesara/link/jax/dispatch/test_subtensor.py b/aesara/link/jax/dispatch/test_subtensor.py index 94d1244f24..03f767976c 100644 --- a/aesara/link/jax/dispatch/test_subtensor.py +++ b/aesara/link/jax/dispatch/test_subtensor.py @@ -1,8 +1,6 @@ -import jax import numpy as np import pytest from jax._src.errors import NonConcreteBooleanIndexError -from packaging.version import parse as version_parse import aesara.tensor as at from aesara.configdefaults import config @@ -11,7 +9,7 @@ from tests.link.jax.test_basic import compare_jax_and_py -def test_jax_Subtensors(): +def test_jax_Subtensor_constant(): # Basic indices x_at = at.as_tensor(np.arange(3 * 4 * 5).reshape((3, 4, 5))) out_at = x_at[1, 2, 0] @@ -19,6 +17,16 @@ def test_jax_Subtensors(): out_fg = FunctionGraph([], [out_at]) compare_jax_and_py(out_fg, []) + out_at = x_at[1:, 1, :] + assert isinstance(out_at.owner.op, at_subtensor.Subtensor) + out_fg = FunctionGraph([], [out_at]) + compare_jax_and_py(out_fg, []) + + out_at = x_at[:2, 1, :] + assert isinstance(out_at.owner.op, at_subtensor.Subtensor) + out_fg = FunctionGraph([], [out_at]) + compare_jax_and_py(out_fg, []) + out_at = x_at[1:2, 1, :] assert isinstance(out_at.owner.op, at_subtensor.Subtensor) out_fg = FunctionGraph([], [out_at]) @@ -46,6 +54,21 @@ def test_jax_Subtensors(): out_fg = FunctionGraph([], [out_at]) compare_jax_and_py(out_fg, []) + # Flipping + out_at = x_at[::-1] + out_fg = FunctionGraph([], [out_at]) + compare_jax_and_py(out_fg, []) + + +@pytest.mark.xfail(reason="`a` should be specified as static when JIT-compiling") +def test_jax_Subtensor_dynamic(): + a = at.iscalar("a") + x = at.arange(3) + out_at = x[:a] + assert isinstance(out_at.owner.op, at_subtensor.Subtensor) + out_fg = FunctionGraph([a], [out_at]) + compare_jax_and_py(out_fg, [1]) + def test_jax_Subtensor_boolean_mask(): """JAX does not support resizing arrays with boolean masks.""" @@ -53,7 +76,7 @@ def test_jax_Subtensor_boolean_mask(): out_at = x_at[x_at < 0] assert isinstance(out_at.owner.op, at_subtensor.AdvancedSubtensor) - with pytest.raises(NotImplementedError): + with pytest.raises(NotImplementedError, match="resizing arrays with boolean"): out_fg = FunctionGraph([], [out_at]) compare_jax_and_py(out_fg, []) From c01a0f394ea56e66361a4a3ac70289ff199702c9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?R=C3=A9mi=20Louf?= Date: Wed, 7 Dec 2022 17:27:38 +0100 Subject: [PATCH 04/14] Use `jax.numpy.copy` directly --- aesara/link/jax/dispatch/basic.py | 18 +----------------- aesara/link/jax/dispatch/elemwise.py | 4 ++-- 2 files changed, 3 insertions(+), 19 deletions(-) diff --git a/aesara/link/jax/dispatch/basic.py b/aesara/link/jax/dispatch/basic.py index 3817389766..f48ea2db1e 100644 --- a/aesara/link/jax/dispatch/basic.py +++ b/aesara/link/jax/dispatch/basic.py @@ -82,26 +82,10 @@ def assert_fn(x, *inputs): return assert_fn -def jnp_safe_copy(x): - try: - res = jnp.copy(x) - except NotImplementedError: - warnings.warn( - "`jnp.copy` is not implemented yet. " "Using the object's `copy` method." - ) - if hasattr(x, "copy"): - res = jnp.array(x.copy()) - else: - warnings.warn(f"Object has no `copy` method: {x}") - res = x - - return res - - @jax_funcify.register(DeepCopyOp) def jax_funcify_DeepCopyOp(op, **kwargs): def deepcopyop(x): - return jnp_safe_copy(x) + return jnp.copy(x) return deepcopyop diff --git a/aesara/link/jax/dispatch/elemwise.py b/aesara/link/jax/dispatch/elemwise.py index b3c4f15be2..5b054fd5fe 100644 --- a/aesara/link/jax/dispatch/elemwise.py +++ b/aesara/link/jax/dispatch/elemwise.py @@ -1,7 +1,7 @@ import jax import jax.numpy as jnp -from aesara.link.jax.dispatch.basic import jax_funcify, jnp_safe_copy +from aesara.link.jax.dispatch.basic import jax_funcify from aesara.tensor.elemwise import CAReduce, DimShuffle, Elemwise from aesara.tensor.special import LogSoftmax, Softmax, SoftmaxGrad @@ -69,7 +69,7 @@ def dimshuffle(x): res = jnp.reshape(res, shape) if not op.inplace: - res = jnp_safe_copy(res) + res = jnp.copy(res) return res From 173e49a572cf2cae3b51080e0f55ece7f346f6b8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?R=C3=A9mi=20Louf?= Date: Wed, 7 Dec 2022 17:39:10 +0100 Subject: [PATCH 05/14] Move `test_subtensor` back to the test suite --- {aesara/link/jax/dispatch => tests/link/jax}/test_subtensor.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename {aesara/link/jax/dispatch => tests/link/jax}/test_subtensor.py (100%) diff --git a/aesara/link/jax/dispatch/test_subtensor.py b/tests/link/jax/test_subtensor.py similarity index 100% rename from aesara/link/jax/dispatch/test_subtensor.py rename to tests/link/jax/test_subtensor.py From 7fb2ba16c2232aba31d42ffe402424d8f251e0bd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?R=C3=A9mi=20Louf?= Date: Thu, 8 Dec 2022 09:25:23 +0100 Subject: [PATCH 06/14] Document discrepancy with `Clip` and `jax.numpy.clip` --- aesara/link/jax/dispatch/scalar.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/aesara/link/jax/dispatch/scalar.py b/aesara/link/jax/dispatch/scalar.py index f169b1e149..e7f627aef2 100644 --- a/aesara/link/jax/dispatch/scalar.py +++ b/aesara/link/jax/dispatch/scalar.py @@ -56,6 +56,14 @@ def identity(x): @jax_funcify.register(Clip) def jax_funcify_Clip(op, **kwargs): + """Register the translation for the `Clip` `Op`. + + Aesara's `Clip` operator operates differently from NumPy's when the + specified `min` is larger than the `max` so we cannot reuse `jax.numpy.clip` + to maintain consistency with Aesara. + + """ + def clip(x, min, max): return jnp.where(x < min, min, jnp.where(x > max, max, x)) From 81816858b712e321b0489f312d38b6a148a01bbe Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?R=C3=A9mi=20Louf?= Date: Thu, 8 Dec 2022 15:17:05 +0100 Subject: [PATCH 07/14] Dispatch some `Op`s to Python operators when scalar inputs --- aesara/link/jax/dispatch/scalar.py | 95 +++++++++++++++++++++++++++++- tests/link/jax/test_scalar.py | 36 +++++++++++ 2 files changed, 129 insertions(+), 2 deletions(-) diff --git a/aesara/link/jax/dispatch/scalar.py b/aesara/link/jax/dispatch/scalar.py index e7f627aef2..a32d906ba9 100644 --- a/aesara/link/jax/dispatch/scalar.py +++ b/aesara/link/jax/dispatch/scalar.py @@ -5,14 +5,57 @@ from aesara.link.jax.dispatch.basic import jax_funcify from aesara.scalar import Softplus -from aesara.scalar.basic import Cast, Clip, Composite, Identity, ScalarOp, Second +from aesara.scalar.basic import ( + Add, + Cast, + Clip, + Composite, + Identity, + IntDiv, + Mod, + Mul, + ScalarOp, + Second, + Sub, +) from aesara.scalar.math import Erf, Erfc, Erfinv, Log1mexp, Psi +def check_if_inputs_scalars(node): + """Check whether all the inputs of an `Elemwise` are scalar values. + + `jax.lax` or `jax.numpy` functions systematically return `TracedArrays`, + while the corresponding Python operators return concrete values when passed + concrete values. In order to be able to compile the largest number of graphs + possible we need to preserve concrete values whenever we can. We thus need + to dispatch differently the Aesara operators depending on whether the inputs + are scalars. + + """ + ndims_input = [inp.type.ndim for inp in node.inputs] + are_inputs_scalars = True + for ndim in ndims_input: + try: + if ndim > 0: + are_inputs_scalars = False + except TypeError: + are_inputs_scalars = False + + return are_inputs_scalars + + @jax_funcify.register(ScalarOp) -def jax_funcify_ScalarOp(op, **kwargs): +def jax_funcify_ScalarOp(op, node, **kwargs): func_name = op.nfunc_spec[0] + # We dispatch some Aesara operators to Python operators + # whenever the inputs are all scalars. + are_inputs_scalars = check_if_inputs_scalars(node) + if are_inputs_scalars: + elemwise = elemwise_scalar(op) + if elemwise is not None: + return elemwise + if "." in func_name: jnp_func = functools.reduce(getattr, [jax] + func_name.split(".")) else: @@ -38,6 +81,54 @@ def elemwise(*args): return jnp_func +@functools.singledispatch +def elemwise_scalar(op): + return None + + +@elemwise_scalar.register(Add) +def elemwise_scalar_add(op): + def elemwise(*inputs): + return sum(inputs) + + return elemwise + + +@elemwise_scalar.register(Mul) +def elemwise_scalar_mul(op): + import operator + from functools import reduce + + def elemwise(*inputs): + return reduce(operator.mul, inputs, 1) + + return elemwise + + +@elemwise_scalar.register(Sub) +def elemwise_scalar_sub(op): + def elemwise(x, y): + return x - y + + return elemwise + + +@elemwise_scalar.register(IntDiv) +def elemwise_scalar_intdiv(op): + def elemwise(x, y): + return x // y + + return elemwise + + +@elemwise_scalar.register(Mod) +def elemwise_scalar_mod(op): + def elemwise(x, y): + return x % y + + return elemwise + + @jax_funcify.register(Cast) def jax_funcify_Cast(op, **kwargs): def cast(x): diff --git a/tests/link/jax/test_scalar.py b/tests/link/jax/test_scalar.py index 02ecf178a7..8d0e107dc2 100644 --- a/tests/link/jax/test_scalar.py +++ b/tests/link/jax/test_scalar.py @@ -151,6 +151,42 @@ def test_jax_variadic_Scalar(): compare_jax_and_py(fgraph, [get_test_value(i) for i in fgraph.inputs]) +def test_add_scalars(): + x = at.matrix("x") + size = x.shape[0] + x.shape[0] + x.shape[1] + out = at.ones(size).astype(config.floatX) + + out_fg = FunctionGraph([x], [out]) + compare_jax_and_py(out_fg, [np.ones((2, 3)).astype(config.floatX)]) + + +def test_mul_scalars(): + x = at.matrix("x") + size = x.shape[0] * x.shape[0] * x.shape[1] + out = at.ones(size).astype(config.floatX) + + out_fg = FunctionGraph([x], [out]) + compare_jax_and_py(out_fg, [np.ones((2, 3)).astype(config.floatX)]) + + +def test_div_scalars(): + x = at.matrix("x") + size = x.shape[0] // x.shape[1] + out = at.ones(size).astype(config.floatX) + + out_fg = FunctionGraph([x], [out]) + compare_jax_and_py(out_fg, [np.ones((12, 3)).astype(config.floatX)]) + + +def test_mod_scalars(): + x = at.matrix("x") + size = x.shape[0] % x.shape[1] + out = at.ones(size).astype(config.floatX) + + out_fg = FunctionGraph([x], [out]) + compare_jax_and_py(out_fg, [np.ones((12, 3)).astype(config.floatX)]) + + def test_jax_multioutput(): x = vector("x") x.tag.test_value = np.r_[1.0, 2.0].astype(config.floatX) From 3ddde9454acfe6b9b9923cdc0ea476585c84a920 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?R=C3=A9mi=20Louf?= Date: Thu, 8 Dec 2022 16:14:54 +0100 Subject: [PATCH 08/14] Implement `TensorFromScalar` as a pass-through --- aesara/link/jax/dispatch/tensor_basic.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/aesara/link/jax/dispatch/tensor_basic.py b/aesara/link/jax/dispatch/tensor_basic.py index a578229fd7..feecf3922a 100644 --- a/aesara/link/jax/dispatch/tensor_basic.py +++ b/aesara/link/jax/dispatch/tensor_basic.py @@ -128,7 +128,7 @@ def makevector(*x): @jax_funcify.register(TensorFromScalar) def jax_funcify_TensorFromScalar(op, **kwargs): def tensor_from_scalar(x): - return jnp.array(x) + return x return tensor_from_scalar From 60fbdad5ca50d76ab92bbe234322cdf9d685f357 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?R=C3=A9mi=20Louf?= Date: Thu, 8 Dec 2022 16:15:26 +0100 Subject: [PATCH 09/14] Simplify the `IncSubtensor` dispatcher --- aesara/link/jax/dispatch/subtensor.py | 38 +++++++++------------------ tests/link/jax/test_subtensor.py | 33 ++++++++++++----------- 2 files changed, 29 insertions(+), 42 deletions(-) diff --git a/aesara/link/jax/dispatch/subtensor.py b/aesara/link/jax/dispatch/subtensor.py index 46293e2300..c1cc8ca209 100644 --- a/aesara/link/jax/dispatch/subtensor.py +++ b/aesara/link/jax/dispatch/subtensor.py @@ -1,5 +1,3 @@ -import jax - from aesara.link.jax.dispatch.basic import jax_funcify from aesara.tensor.subtensor import ( AdvancedIncSubtensor, @@ -33,7 +31,7 @@ """ -def assert_indices_jax_compatible(node, idx_list): +def subtensor_assert_indices_jax_compatible(node, idx_list): from aesara.graph.basic import Constant from aesara.tensor.var import TensorVariable @@ -55,7 +53,7 @@ def assert_indices_jax_compatible(node, idx_list): def jax_funcify_Subtensor(op, node, **kwargs): idx_list = getattr(op, "idx_list", None) - assert_indices_jax_compatible(node, idx_list) + subtensor_assert_indices_jax_compatible(node, idx_list) def subtensor_constant(x, *ilists): indices = indices_from_subtensor(ilists, idx_list) @@ -69,25 +67,19 @@ def subtensor_constant(x, *ilists): @jax_funcify.register(IncSubtensor) @jax_funcify.register(AdvancedIncSubtensor1) -def jax_funcify_IncSubtensor(op, **kwargs): +def jax_funcify_IncSubtensor(op, node, **kwargs): idx_list = getattr(op, "idx_list", None) if getattr(op, "set_instead_of_inc", False): - jax_fn = getattr(jax.ops, "index_update", None) - - if jax_fn is None: - def jax_fn(x, indices, y): - return x.at[indices].set(y) + def jax_fn(x, indices, y): + return x.at[indices].set(y) else: - jax_fn = getattr(jax.ops, "index_add", None) - - if jax_fn is None: - def jax_fn(x, indices, y): - return x.at[indices].add(y) + def jax_fn(x, indices, y): + return x.at[indices].add(y) def incsubtensor(x, y, *ilist, jax_fn=jax_fn, idx_list=idx_list): indices = indices_from_subtensor(ilist, idx_list) @@ -100,23 +92,17 @@ def incsubtensor(x, y, *ilist, jax_fn=jax_fn, idx_list=idx_list): @jax_funcify.register(AdvancedIncSubtensor) -def jax_funcify_AdvancedIncSubtensor(op, **kwargs): +def jax_funcify_AdvancedIncSubtensor(op, node, **kwargs): if getattr(op, "set_instead_of_inc", False): - jax_fn = getattr(jax.ops, "index_update", None) - if jax_fn is None: - - def jax_fn(x, indices, y): - return x.at[indices].set(y) + def jax_fn(x, indices, y): + return x.at[indices].set(y) else: - jax_fn = getattr(jax.ops, "index_add", None) - - if jax_fn is None: - def jax_fn(x, indices, y): - return x.at[indices].add(y) + def jax_fn(x, indices, y): + return x.at[indices].add(y) def advancedincsubtensor(x, y, *ilist, jax_fn=jax_fn): return jax_fn(x, ilist, y) diff --git a/tests/link/jax/test_subtensor.py b/tests/link/jax/test_subtensor.py index 03f767976c..c3a6aa1d99 100644 --- a/tests/link/jax/test_subtensor.py +++ b/tests/link/jax/test_subtensor.py @@ -1,6 +1,5 @@ import numpy as np import pytest -from jax._src.errors import NonConcreteBooleanIndexError import aesara.tensor as at from aesara.configdefaults import config @@ -179,7 +178,11 @@ def test_jax_IncSubtensor(): compare_jax_and_py(out_fg, []) -def test_jax_IncSubtensors_unsupported(): +@pytest.mark.xfail( + reason="Re-expressible boolean logic. We need a rewrite Aesara-side to remove the DimShuffle." +) +def test_jax_IncSubtensor_boolean_mask_reexpressible(): + """Some boolean logic can be re-expressed and JIT-compiled""" rng = np.random.default_rng(213234) x_np = rng.uniform(-1, 1, size=(3, 4, 5)).astype(config.floatX) x_at = at.constant(np.arange(3 * 4 * 5).reshape((3, 4, 5)).astype(config.floatX)) @@ -188,30 +191,28 @@ def test_jax_IncSubtensors_unsupported(): out_at = at_subtensor.set_subtensor(x_at[mask_at], 0.0) assert isinstance(out_at.owner.op, at_subtensor.AdvancedIncSubtensor) out_fg = FunctionGraph([], [out_at]) - with pytest.raises( - NonConcreteBooleanIndexError, match="Array boolean indices must be concrete" - ): - compare_jax_and_py(out_fg, []) + compare_jax_and_py(out_fg, []) - mask_at = at.as_tensor_variable(x_np) > 0 - out_at = at_subtensor.set_subtensor(x_at[mask_at], 1.0) + mask_at = at.as_tensor(x_np) > 0 + out_at = at_subtensor.inc_subtensor(x_at[mask_at], 1.0) assert isinstance(out_at.owner.op, at_subtensor.AdvancedIncSubtensor) out_fg = FunctionGraph([], [out_at]) - with pytest.raises( - NonConcreteBooleanIndexError, match="Array boolean indices must be concrete" - ): - compare_jax_and_py(out_fg, []) + compare_jax_and_py(out_fg, []) + + +def test_jax_IncSubtensors_unsupported(): + rng = np.random.default_rng(213234) + x_np = rng.uniform(-1, 1, size=(3, 4, 5)).astype(config.floatX) + x_at = at.constant(np.arange(3 * 4 * 5).reshape((3, 4, 5)).astype(config.floatX)) st_at = at.as_tensor_variable(x_np[[0, 2], 0, :3]) out_at = at_subtensor.set_subtensor(x_at[[0, 2], 0, :3], st_at) assert isinstance(out_at.owner.op, at_subtensor.AdvancedIncSubtensor) out_fg = FunctionGraph([], [out_at]) - with pytest.raises(IndexError, match="Array slice indices must have static"): - compare_jax_and_py(out_fg, []) + compare_jax_and_py(out_fg, []) st_at = at.as_tensor_variable(x_np[[0, 2], 0, :3]) out_at = at_subtensor.inc_subtensor(x_at[[0, 2], 0, :3], st_at) assert isinstance(out_at.owner.op, at_subtensor.AdvancedIncSubtensor) out_fg = FunctionGraph([], [out_at]) - with pytest.raises(IndexError, match="Array slice indices must have static"): - compare_jax_and_py(out_fg, []) + compare_jax_and_py(out_fg, []) From f9159e7940c445a930dee2c8e6c4fc027f8b8fc1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?R=C3=A9mi=20Louf?= Date: Mon, 12 Dec 2022 14:42:01 +0100 Subject: [PATCH 10/14] Add rewrites to re-express boolean indexing logic --- aesara/tensor/rewriting/__init__.py | 3 ++ aesara/tensor/rewriting/jax.py | 78 +++++++++++++++++++++++++++++ tests/link/jax/test_subtensor.py | 70 ++++++++++++++------------ 3 files changed, 119 insertions(+), 32 deletions(-) create mode 100644 aesara/tensor/rewriting/jax.py diff --git a/aesara/tensor/rewriting/__init__.py b/aesara/tensor/rewriting/__init__.py index d1ad5b3c51..3b4421b439 100644 --- a/aesara/tensor/rewriting/__init__.py +++ b/aesara/tensor/rewriting/__init__.py @@ -1,6 +1,9 @@ import aesara.tensor.rewriting.basic import aesara.tensor.rewriting.elemwise import aesara.tensor.rewriting.extra_ops + +# Register JAX specializations +import aesara.tensor.rewriting.jax import aesara.tensor.rewriting.math import aesara.tensor.rewriting.shape import aesara.tensor.rewriting.special diff --git a/aesara/tensor/rewriting/jax.py b/aesara/tensor/rewriting/jax.py new file mode 100644 index 0000000000..8e57755c16 --- /dev/null +++ b/aesara/tensor/rewriting/jax.py @@ -0,0 +1,78 @@ +from aesara.compile import optdb +from aesara.graph.rewriting.basic import in2out, node_rewriter +from aesara.tensor.var import TensorVariable +import aesara.tensor as at +from aesara.tensor.subtensor import AdvancedIncSubtensor, AdvancedSubtensor +from aesara.tensor.math import Sum + + +@node_rewriter([AdvancedIncSubtensor]) +def boolean_indexing_set_or_inc(fgraph, node): + """Replace `AdvancedIncSubtensor` when using boolean indexing using `Switch`. + + JAX cannot JIT-compile functions that use boolean indexing to set values in + an array. A workaround is to re-express this logic using `jax.numpy.where`. + This rewrite allows to improve upon JAX's API. + + """ + + op = node.op + x = node.inputs[0] + y = node.inputs[1] + cond = node.inputs[2] + + if not isinstance(cond, TensorVariable): + return + + if not cond.type.dtype == 'bool': + return + + if op.set_instead_of_inc: + out = at.where(cond, y, x) + return out.owner.outputs + else: + out = at.where(cond, x + y, x) + return out.owner.outputs + + +optdb.register( + "jax_boolean_indexing_set_or_inc", in2out(boolean_indexing_set_or_inc), "jax", position=100 +) + + +@node_rewriter([Sum]) +def boolean_indexing_sum(fgraph, node): + """Replace the sum of `AdvancedSubtensor` with boolean indexing. + + JAX cannot JIT-compile functions that use boolean indexing, but can compile + those expressions that can be re-expressed using `jax.numpy.where`. This + rewrite re-rexpressed the model on the behalf of the user and thus allows to + improve upon JAX's API. + + """ + operand = node.inputs[0] + + if not isinstance(operand, TensorVariable): + return + + if operand.owner is None: + return + + if not isinstance(operand.owner.op, AdvancedSubtensor): + return + + x = operand.owner.inputs[0] + cond = operand.owner.inputs[1] + + if not isinstance(cond, TensorVariable): + return + + if not cond.type.dtype == 'bool': + return + + out = at.sum(at.where(cond, x, 0)) + return out.owner.outputs + +optdb.register( + "jax_boolean_indexing_sum", in2out(boolean_indexing_sum), "jax", position=100 +) diff --git a/tests/link/jax/test_subtensor.py b/tests/link/jax/test_subtensor.py index c3a6aa1d99..fba311c003 100644 --- a/tests/link/jax/test_subtensor.py +++ b/tests/link/jax/test_subtensor.py @@ -80,15 +80,21 @@ def test_jax_Subtensor_boolean_mask(): compare_jax_and_py(out_fg, []) -@pytest.mark.xfail( - reason="Re-expressible boolean logic. We need a rewrite Aesara-side." -) def test_jax_Subtensor_boolean_mask_reexpressible(): - """Some boolean logic can be re-expressed and JIT-compiled""" - x_at = at.arange(-5, 5) + """Summing values with boolean indexing. + + This test ensures that the sum of an `AdvancedSubtensor` `Op`s with boolean + indexing is replaced with the sum of an equivalent `Switch` `Op`, using the + `jax_boolean_indexing_sum` rewrite. + + JAX forces users to re-express this logic manually, so this is an + improvement over its user interface. + + """ + x_at = at.vector("x") out_at = x_at[x_at < 0].sum() - out_fg = FunctionGraph([], [out_at]) - compare_jax_and_py(out_fg, []) + out_fg = FunctionGraph([x_at], [out_at]) + compare_jax_and_py(out_fg, [np.arange(-5, 5).astype(config.floatX)]) def test_jax_IncSubtensor(): @@ -177,42 +183,42 @@ def test_jax_IncSubtensor(): out_fg = FunctionGraph([], [out_at]) compare_jax_and_py(out_fg, []) - -@pytest.mark.xfail( - reason="Re-expressible boolean logic. We need a rewrite Aesara-side to remove the DimShuffle." -) -def test_jax_IncSubtensor_boolean_mask_reexpressible(): - """Some boolean logic can be re-expressed and JIT-compiled""" - rng = np.random.default_rng(213234) - x_np = rng.uniform(-1, 1, size=(3, 4, 5)).astype(config.floatX) - x_at = at.constant(np.arange(3 * 4 * 5).reshape((3, 4, 5)).astype(config.floatX)) - - mask_at = at.as_tensor(x_np) > 0 - out_at = at_subtensor.set_subtensor(x_at[mask_at], 0.0) + st_at = at.as_tensor_variable(x_np[[0, 2], 0, :3]) + out_at = at_subtensor.set_subtensor(x_at[[0, 2], 0, :3], st_at) assert isinstance(out_at.owner.op, at_subtensor.AdvancedIncSubtensor) out_fg = FunctionGraph([], [out_at]) compare_jax_and_py(out_fg, []) - mask_at = at.as_tensor(x_np) > 0 - out_at = at_subtensor.inc_subtensor(x_at[mask_at], 1.0) + st_at = at.as_tensor_variable(x_np[[0, 2], 0, :3]) + out_at = at_subtensor.inc_subtensor(x_at[[0, 2], 0, :3], st_at) assert isinstance(out_at.owner.op, at_subtensor.AdvancedIncSubtensor) out_fg = FunctionGraph([], [out_at]) compare_jax_and_py(out_fg, []) -def test_jax_IncSubtensors_unsupported(): +def test_jax_IncSubtensor_boolean_indexing_reexpressible(): + """Setting or incrementing values with boolean indexing. + + This test ensures that `AdvancedIncSubtensor` `Op`s with boolean indexing is + replaced with an equivalent `Switch` `Op`, using the + `jax_boolean_indexing_set_of_inc` rewrite. + + JAX forces users to re-express this logic manually, so this is an + improvement over its user interface. + + """ rng = np.random.default_rng(213234) - x_np = rng.uniform(-1, 1, size=(3, 4, 5)).astype(config.floatX) - x_at = at.constant(np.arange(3 * 4 * 5).reshape((3, 4, 5)).astype(config.floatX)) + x_np = rng.uniform(-1, 1, size=(4, 5)).astype(config.floatX) - st_at = at.as_tensor_variable(x_np[[0, 2], 0, :3]) - out_at = at_subtensor.set_subtensor(x_at[[0, 2], 0, :3], st_at) + x_at = at.matrix("x") + mask_at = at.as_tensor(x_at) > 0 + out_at = at_subtensor.set_subtensor(x_at[mask_at], 0.0) assert isinstance(out_at.owner.op, at_subtensor.AdvancedIncSubtensor) - out_fg = FunctionGraph([], [out_at]) - compare_jax_and_py(out_fg, []) + out_fg = FunctionGraph([x_at], [out_at]) + compare_jax_and_py(out_fg, [x_np]) - st_at = at.as_tensor_variable(x_np[[0, 2], 0, :3]) - out_at = at_subtensor.inc_subtensor(x_at[[0, 2], 0, :3], st_at) + mask_at = at.as_tensor(x_at) > 0 + out_at = at_subtensor.inc_subtensor(x_at[mask_at], 1.0) assert isinstance(out_at.owner.op, at_subtensor.AdvancedIncSubtensor) - out_fg = FunctionGraph([], [out_at]) - compare_jax_and_py(out_fg, []) + out_fg = FunctionGraph([x_at], [out_at]) + compare_jax_and_py(out_fg, [x_np]) From 9cc0c286ff0e4d5a47ea6d49843f1d7467192b7d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?R=C3=A9mi=20Louf?= Date: Mon, 12 Dec 2022 15:41:24 +0100 Subject: [PATCH 11/14] Typify 0-dim arrays to corresponding number --- aesara/link/jax/dispatch/basic.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/aesara/link/jax/dispatch/basic.py b/aesara/link/jax/dispatch/basic.py index f48ea2db1e..c24eb8329f 100644 --- a/aesara/link/jax/dispatch/basic.py +++ b/aesara/link/jax/dispatch/basic.py @@ -30,6 +30,8 @@ def jax_typify(data, dtype=None, **kwargs): @jax_typify.register(np.ndarray) def jax_typify_ndarray(data, dtype=None, **kwargs): + if len(data.shape) == 0: + return data.item() return jnp.array(data, dtype=dtype) From 0202f6b901fbdfbe68632e450f2da50b389a6660 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?R=C3=A9mi=20Louf?= Date: Mon, 12 Dec 2022 15:41:54 +0100 Subject: [PATCH 12/14] Refactor the JAX implementation of `Reshape` --- aesara/link/jax/dispatch/shape.py | 30 +++++++++++++- aesara/tensor/rewriting/jax.py | 69 ++++++++++++++++++++++++++++--- tests/link/jax/test_shape.py | 38 +++++++++-------- 3 files changed, 113 insertions(+), 24 deletions(-) diff --git a/aesara/link/jax/dispatch/shape.py b/aesara/link/jax/dispatch/shape.py index c9a12b7b57..df156fa7c2 100644 --- a/aesara/link/jax/dispatch/shape.py +++ b/aesara/link/jax/dispatch/shape.py @@ -28,11 +28,38 @@ def shape_tuple_fn(*x): return shape_tuple_fn +SHAPE_NOT_COMPATIBLE = """JAX requires concrete values for the `shape` parameter of `jax.numpy.reshape`. +Concrete values are either constants: + +>>> import aesara.tensor as at +>>> x = at.ones(6) +>>> y = x.reshape((2, 3)) + +Or the shape of an array: + +>>> mat = at.matrix('mat') +>>> y = x.reshape(mat.shape) +""" + + +def assert_shape_argument_jax_compatible(shape): + """Assert whether the current node can be JIT-compiled by JAX. + + JAX can JIT-compile functions with a `shape` or `size` argument if it is + given a concrete value, i.e. either a constant or the shape of any traced + value. + + """ + shape_op = shape.owner.op + if not isinstance(shape_op, (Shape, Shape_i, JAXShapeTuple)): + raise NotImplementedError(SHAPE_NOT_COMPATIBLE) + + @jax_funcify.register(Reshape) def jax_funcify_Reshape(op, node, **kwargs): - # JAX reshape only works with constant inputs, otherwise JIT fails shape = node.inputs[1] + if isinstance(shape, Constant): constant_shape = shape.data @@ -40,6 +67,7 @@ def reshape(x, shape): return jnp.reshape(x, constant_shape) else: + assert_shape_argument_jax_compatible(shape) def reshape(x, shape): return jnp.reshape(x, shape) diff --git a/aesara/tensor/rewriting/jax.py b/aesara/tensor/rewriting/jax.py index 8e57755c16..6008512a92 100644 --- a/aesara/tensor/rewriting/jax.py +++ b/aesara/tensor/rewriting/jax.py @@ -1,9 +1,12 @@ +import aesara.tensor as at from aesara.compile import optdb from aesara.graph.rewriting.basic import in2out, node_rewriter -from aesara.tensor.var import TensorVariable -import aesara.tensor as at -from aesara.tensor.subtensor import AdvancedIncSubtensor, AdvancedSubtensor +from aesara.tensor.basic import MakeVector +from aesara.tensor.elemwise import DimShuffle from aesara.tensor.math import Sum +from aesara.tensor.shape import Reshape +from aesara.tensor.subtensor import AdvancedIncSubtensor, AdvancedSubtensor +from aesara.tensor.var import TensorVariable @node_rewriter([AdvancedIncSubtensor]) @@ -24,7 +27,7 @@ def boolean_indexing_set_or_inc(fgraph, node): if not isinstance(cond, TensorVariable): return - if not cond.type.dtype == 'bool': + if not cond.type.dtype == "bool": return if op.set_instead_of_inc: @@ -36,7 +39,10 @@ def boolean_indexing_set_or_inc(fgraph, node): optdb.register( - "jax_boolean_indexing_set_or_inc", in2out(boolean_indexing_set_or_inc), "jax", position=100 + "jax_boolean_indexing_set_or_inc", + in2out(boolean_indexing_set_or_inc), + "jax", + position=100, ) @@ -67,12 +73,63 @@ def boolean_indexing_sum(fgraph, node): if not isinstance(cond, TensorVariable): return - if not cond.type.dtype == 'bool': + if not cond.type.dtype == "bool": return out = at.sum(at.where(cond, x, 0)) return out.owner.outputs + optdb.register( "jax_boolean_indexing_sum", in2out(boolean_indexing_sum), "jax", position=100 ) + + +@node_rewriter([Reshape]) +def shape_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 `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 aesara.link.jax.dispatch.shape import JAXShapeTuple + + shape_arg = node.inputs[1] + shape_node = shape_arg.owner + + if shape_node is None: + return + + if isinstance(shape_node.op, JAXShapeTuple): + return + + if isinstance(shape_node.op, MakeVector) or ( + isinstance(shape_node.op, DimShuffle) + and shape_node.op.input_broadcastable == () + and shape_node.op.new_order == ("x",) + ): + # Here Aesara converted a tuple or list to a tensor + new_shape_args = JAXShapeTuple()(*shape_node.inputs) + new_inputs = list(node.inputs) + new_inputs[1] = new_shape_args + + new_node = node.clone_with_new_inputs(new_inputs) + return new_node.outputs + + +optdb.register( + "jax_shape_parameter_as_tuple", + in2out(shape_parameter_as_tuple), + "jax", + position=100, +) diff --git a/tests/link/jax/test_shape.py b/tests/link/jax/test_shape.py index 6b1bd442fa..c9f107f384 100644 --- a/tests/link/jax/test_shape.py +++ b/tests/link/jax/test_shape.py @@ -45,30 +45,34 @@ def test_jax_specify_shape(): compare_jax_and_py(x_fg, []) -def test_jax_Reshape(): +def test_jax_Reshape_constant(): a = vector("a") x = reshape(a, (2, 2)) x_fg = FunctionGraph([a], [x]) compare_jax_and_py(x_fg, [np.r_[1.0, 2.0, 3.0, 4.0].astype(config.floatX)]) - # Test breaking "omnistaging" changes in JAX. - # See https://github.com/tensorflow/probability/commit/782d0c64eb774b9aac54a1c8488e4f1f96fbbc68 + +def test_jax_Reshape_concrete_shape(): + """JAX should compile when a concrete value is passed for the `shape` parameter.""" + a = vector("a") + x = reshape(a, a.shape) + x_fg = FunctionGraph([a], [x]) + compare_jax_and_py(x_fg, [np.r_[1.0, 2.0, 3.0, 4.0].astype(config.floatX)]) + x = reshape(a, (a.shape[0] // 2, a.shape[0] // 2)) x_fg = FunctionGraph([a], [x]) - with pytest.raises( - TypeError, - match="Shapes must be 1D sequences of concrete values of integer type", - ): - compare_jax_and_py(x_fg, [np.r_[1.0, 2.0, 3.0, 4.0].astype(config.floatX)]) - - b = iscalar("b") - x = reshape(a, (b, b)) - x_fg = FunctionGraph([a, b], [x]) - with pytest.raises( - TypeError, - match="Shapes must be 1D sequences of concrete values of integer type", - ): - compare_jax_and_py(x_fg, [np.r_[1.0, 2.0, 3.0, 4.0].astype(config.floatX), 2]) + compare_jax_and_py(x_fg, [np.r_[1.0, 2.0, 3.0, 4.0].astype(config.floatX)]) + + +@pytest.mark.xfail( + reason="`shape_at` should be specified as a static argument", strict=True +) +def test_jax_Reshape_shape_graph_input(): + a = vector("a") + shape_at = iscalar("b") + x = reshape(a, (shape_at, shape_at)) + x_fg = FunctionGraph([a, shape_at], [x]) + compare_jax_and_py(x_fg, [np.r_[1.0, 2.0, 3.0, 4.0].astype(config.floatX), 2]) def test_jax_compile_ops(): From e5f959eb7dfe09f8c39bde951cacb4b399c8ae07 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?R=C3=A9mi=20Louf?= Date: Mon, 12 Dec 2022 16:10:55 +0100 Subject: [PATCH 13/14] Fix the JAX implementation of `SpecifyShape` --- aesara/link/jax/dispatch/shape.py | 4 ++-- tests/link/jax/test_shape.py | 35 +++++++++++++------------------ 2 files changed, 17 insertions(+), 22 deletions(-) diff --git a/aesara/link/jax/dispatch/shape.py b/aesara/link/jax/dispatch/shape.py index df156fa7c2..6fd7a74fb1 100644 --- a/aesara/link/jax/dispatch/shape.py +++ b/aesara/link/jax/dispatch/shape.py @@ -94,10 +94,10 @@ def shape_i(x): @jax_funcify.register(SpecifyShape) -def jax_funcify_SpecifyShape(op, **kwargs): +def jax_funcify_SpecifyShape(op, node, **kwargs): def specifyshape(x, *shape): assert x.ndim == len(shape) - assert jnp.all(x.shape == tuple(shape)), ( + assert x.shape == tuple(shape), ( "got shape", x.shape, "expected", diff --git a/tests/link/jax/test_shape.py b/tests/link/jax/test_shape.py index c9f107f384..19e5f7d813 100644 --- a/tests/link/jax/test_shape.py +++ b/tests/link/jax/test_shape.py @@ -1,13 +1,11 @@ -import jax import numpy as np import pytest -from packaging.version import parse as version_parse import aesara.tensor as at from aesara.compile.ops import DeepCopyOp, ViewOp from aesara.configdefaults import config from aesara.graph.fg import FunctionGraph -from aesara.tensor.shape import Shape, Shape_i, SpecifyShape, Unbroadcast, reshape +from aesara.tensor.shape import Shape, Shape_i, Unbroadcast, reshape from aesara.tensor.type import iscalar, vector from tests.link.jax.test_basic import compare_jax_and_py @@ -25,24 +23,21 @@ def test_jax_shape_ops(): compare_jax_and_py(x_fg, [], must_be_device_array=False) -@pytest.mark.xfail( - version_parse(jax.__version__) >= version_parse("0.2.12"), - reason="Omnistaging cannot be disabled", -) def test_jax_specify_shape(): - x_np = np.zeros((20, 3)) - x = SpecifyShape()(at.as_tensor_variable(x_np), (20, 3)) - x_fg = FunctionGraph([], [x]) - - compare_jax_and_py(x_fg, []) - - with config.change_flags(compute_test_value="off"): - - x = SpecifyShape()(at.as_tensor_variable(x_np), *(2, 3)) - x_fg = FunctionGraph([], [x]) - - with pytest.raises(AssertionError): - compare_jax_and_py(x_fg, []) + in_at = at.matrix("in") + x = at.specify_shape(in_at, (4, 5)) + x_fg = FunctionGraph([in_at], [x]) + compare_jax_and_py(x_fg, [np.ones((4, 5)).astype(config.floatX)]) + + # When used to assert two arrays have similar shapes + in_at = at.matrix("in") + shape_at = at.matrix("shape") + x = at.specify_shape(in_at, shape_at.shape) + x_fg = FunctionGraph([in_at, shape_at], [x]) + compare_jax_and_py( + x_fg, + [np.ones((4, 5)).astype(config.floatX), np.ones((4, 5)).astype(config.floatX)], + ) def test_jax_Reshape_constant(): From 78d2dbf3721b9a5ec597d1907dcfb24694692e51 Mon Sep 17 00:00:00 2001 From: "Brandon T. Willard" Date: Mon, 12 Dec 2022 23:06:03 -0600 Subject: [PATCH 14/14] Use separate DB queries for each JAX test mode --- tests/link/jax/test_basic.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/tests/link/jax/test_basic.py b/tests/link/jax/test_basic.py index 71388d72db..0f2d53a919 100644 --- a/tests/link/jax/test_basic.py +++ b/tests/link/jax/test_basic.py @@ -27,9 +27,12 @@ def set_aesara_flags(): jax = pytest.importorskip("jax") -opts = RewriteDatabaseQuery(include=["jax"], exclude=["cxx_only", "BlasOpt"]) -jax_mode = Mode(JAXLinker(), opts) -py_mode = Mode("py", opts) +jax_mode = Mode( + JAXLinker(), RewriteDatabaseQuery(include=["jax"], exclude=["cxx_only", "BlasOpt"]) +) +py_mode = Mode( + "py", RewriteDatabaseQuery(include=[None], exclude=["cxx_only", "BlasOpt"]) +) def compare_jax_and_py(