From 26576273ab8e65c65af67f20dc07ae577c04b306 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?R=C3=A9mi=20Louf?= Date: Fri, 14 Oct 2022 16:50:51 +0200 Subject: [PATCH 1/6] Pass updates to `compare_jax_and_py` --- tests/link/jax/test_basic.py | 39 ++++++++++++++------------------ tests/link/jax/test_extra_ops.py | 4 +--- tests/link/jax/test_shape.py | 4 ++-- 3 files changed, 20 insertions(+), 27 deletions(-) diff --git a/tests/link/jax/test_basic.py b/tests/link/jax/test_basic.py index 0f2d53a919..1a7fa1f6e6 100644 --- a/tests/link/jax/test_basic.py +++ b/tests/link/jax/test_basic.py @@ -37,9 +37,10 @@ def set_aesara_flags(): def compare_jax_and_py( fgraph: FunctionGraph, - test_inputs: Iterable, + inputs: Iterable, assert_fn: Optional[Callable] = None, - must_be_device_array: bool = True, + jax_mode=jax_mode, + updates=None, ): """Function to compare python graph output and jax compiled output for testing equality @@ -56,34 +57,28 @@ def compare_jax_and_py( assert_fn: func, opt Assert function used to check for equality between python and jax. If not provided uses np.testing.assert_allclose - must_be_device_array: Bool - Checks for instance of jax.interpreters.xla.DeviceArray. For testing purposes - if this device array is found it indicates if the result was computed by jax - - Returns - ------- - jax_res + updates + Updates to be passed to `aesara.function`. """ if assert_fn is None: assert_fn = partial(np.testing.assert_allclose, rtol=1e-4) - fn_inputs = [i for i in fgraph.inputs if not isinstance(i, SharedVariable)] - aesara_jax_fn = function(fn_inputs, fgraph.outputs, mode=jax_mode) - jax_res = aesara_jax_fn(*test_inputs) + if isinstance(fgraph, tuple): + fn_inputs, fn_outputs = fgraph + else: + fn_inputs = fgraph.inputs + fn_outputs = fgraph.outputs + + fn_inputs = [i for i in fn_inputs if not isinstance(i, SharedVariable)] - if must_be_device_array: - if isinstance(jax_res, list): - assert all( - isinstance(res, jax.interpreters.xla.DeviceArray) for res in jax_res - ) - else: - assert isinstance(jax_res, jax.interpreters.xla.DeviceArray) + aesara_py_fn = function(fn_inputs, fn_outputs, mode=py_mode, updates=updates) + py_res = aesara_py_fn(*inputs) - aesara_py_fn = function(fn_inputs, fgraph.outputs, mode=py_mode) - py_res = aesara_py_fn(*test_inputs) + aesara_jax_fn = function(fn_inputs, fn_outputs, mode=jax_mode, updates=updates) + jax_res = aesara_jax_fn(*inputs) - if len(fgraph.outputs) > 1: + if len(fn_outputs) > 1: for j, p in zip(jax_res, py_res): assert_fn(j, p) else: diff --git a/tests/link/jax/test_extra_ops.py b/tests/link/jax/test_extra_ops.py index 8c9b70ef37..1221d4952d 100644 --- a/tests/link/jax/test_extra_ops.py +++ b/tests/link/jax/test_extra_ops.py @@ -58,9 +58,7 @@ def test_extra_ops(): indices = np.arange(np.product((3, 4))) out = at_extra_ops.unravel_index(indices, (3, 4), order="C") fgraph = FunctionGraph([], out) - compare_jax_and_py( - fgraph, [get_test_value(i) for i in fgraph.inputs], must_be_device_array=False - ) + compare_jax_and_py(fgraph, [get_test_value(i) for i in fgraph.inputs]) @pytest.mark.parametrize( diff --git a/tests/link/jax/test_shape.py b/tests/link/jax/test_shape.py index 19e5f7d813..ada3b64b3d 100644 --- a/tests/link/jax/test_shape.py +++ b/tests/link/jax/test_shape.py @@ -15,12 +15,12 @@ def test_jax_shape_ops(): x = Shape()(at.as_tensor_variable(x_np)) x_fg = FunctionGraph([], [x]) - compare_jax_and_py(x_fg, [], must_be_device_array=False) + compare_jax_and_py(x_fg, []) x = Shape_i(1)(at.as_tensor_variable(x_np)) x_fg = FunctionGraph([], [x]) - compare_jax_and_py(x_fg, [], must_be_device_array=False) + compare_jax_and_py(x_fg, []) def test_jax_specify_shape(): From 3918d9ff0ad6cd2aa8ab1953af35a47b0d89e9cf Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?R=C3=A9mi=20Louf?= Date: Fri, 14 Oct 2022 16:51:27 +0200 Subject: [PATCH 2/6] Refactor the JAX `Scan` dispatcher --- aesara/link/jax/dispatch/scan.py | 382 ++++++++++++++++++++----------- tests/link/jax/test_scan.py | 332 +++++++++++++++++++++++++++ 2 files changed, 581 insertions(+), 133 deletions(-) diff --git a/aesara/link/jax/dispatch/scan.py b/aesara/link/jax/dispatch/scan.py index 12c588c0d6..554939d1f2 100644 --- a/aesara/link/jax/dispatch/scan.py +++ b/aesara/link/jax/dispatch/scan.py @@ -1,159 +1,275 @@ +from collections import defaultdict +from typing import Callable, Dict, List + import jax -import jax.numpy as jnp -from aesara.graph.fg import FunctionGraph from aesara.link.jax.dispatch.basic import jax_funcify from aesara.scan.op import Scan -from aesara.scan.utils import ScanArgs +from aesara.tensor.var import TensorVariable @jax_funcify.register(Scan) -def jax_funcify_Scan(op, **kwargs): - inner_fg = FunctionGraph(op.inputs, op.outputs) - jax_at_inner_func = jax_funcify(inner_fg, **kwargs) +def jax_funcify_Scan(op, node, **kwargs): + scan_inner_fn = jax_funcify(op.fgraph) + input_taps = { + "mit_sot": op.info.mit_sot_in_slices, + "sit_sot": op.info.sit_sot_in_slices, + } - def scan(*outer_inputs): - scan_args = ScanArgs( - list(outer_inputs), [None] * op.info.n_outs, op.inputs, op.outputs, op.info - ) + # Outer-inputs are the inputs to the `Scan` apply node, built from the + # the variables provided by the caller to the `scan` function at construction + # time. + def parse_outer_inputs(outer_inputs): + outer_in = { + "n_steps": outer_inputs[0], + "sequences": list(op.outer_seqs(outer_inputs)), + "mit_mot": list(op.outer_mitmot(outer_inputs)), + "mit_sot": list(op.outer_mitsot(outer_inputs)), + "nit_sot": list(op.outer_nitsot(outer_inputs)), + "sit_sot": list(op.outer_sitsot(outer_inputs)), + "shared": list(op.outer_shared(outer_inputs)), + "non_sequences": list(op.outer_non_seqs(outer_inputs)), + } + if len(outer_in["mit_mot"]) > 0: + raise NotImplementedError("mit-mot not supported") + + return outer_in - # `outer_inputs` is a list with the following composite form: - # [n_steps] - # + outer_in_seqs - # + outer_in_mit_mot - # + outer_in_mit_sot - # + outer_in_sit_sot - # + outer_in_shared - # + outer_in_nit_sot - # + outer_in_non_seqs - n_steps = scan_args.n_steps - seqs = scan_args.outer_in_seqs - - # TODO: mit_mots - mit_mot_in_slices = [] - - mit_sot_in_slices = [] - for tap, seq in zip(scan_args.mit_sot_in_slices, scan_args.outer_in_mit_sot): - neg_taps = [abs(t) for t in tap if t < 0] - pos_taps = [abs(t) for t in tap if t > 0] - max_neg = max(neg_taps) if neg_taps else 0 - max_pos = max(pos_taps) if pos_taps else 0 - init_slice = seq[: max_neg + max_pos] - mit_sot_in_slices.append(init_slice) - - sit_sot_in_slices = [seq[0] for seq in scan_args.outer_in_sit_sot] - - init_carry = ( - mit_mot_in_slices, - mit_sot_in_slices, - sit_sot_in_slices, - scan_args.outer_in_shared, - scan_args.outer_in_non_seqs, + if op.info.as_while: + raise NotImplementedError("While loops are not supported in the JAX backend.") + else: + return make_jax_scan_fn( + scan_inner_fn, + parse_outer_inputs, + input_taps, ) - def jax_args_to_inner_scan(op, carry, x): - # `carry` contains all inner-output taps, non_seqs, and shared - # terms - ( - inner_in_mit_mot, - inner_in_mit_sot, - inner_in_sit_sot, - inner_in_shared, - inner_in_non_seqs, - ) = carry - - # `x` contains the in_seqs + +def make_jax_scan_fn( + scan_inner_fn: Callable, + parse_outer_inputs: Callable[[TensorVariable], Dict[str, List[TensorVariable]]], + input_taps: Dict, +): + """Create a `jax.lax.scan` function to perform `Scan` computations. + + `jax.lax.scan` takes an initial `carry` value and a sequence it scans over, + or a number of iterations. The first output of the loop body function, the + `carry`, is carried over to the next iteration. The second, the `output`, is + stacked to the previous outputs. We use this to our advantage to build + `Scan` outputs without having to post-process the storage arrays. + + The JAX scan function needs to perform the following operations: + 1. Extract the inner-inputs; + 2. Build the initial `carry` and `sequence` values; + 3. Inside the loop: + 1. `carry` + sequence elements -> inner-inputs + 2. inner-outputs -> `carry` + 3. inner-outputs -> `output` + 4. Append the last `shared` value to the stacked `output`s + + """ + + def build_jax_scan_inputs(outer_in: Dict): + """Build the inputs to `jax.lax.scan` from the outer-inputs.""" + n_steps = outer_in["n_steps"] + sequences = outer_in["sequences"] + init_carry = { + name: outer_in[name] + for name in ["mit_sot", "sit_sot", "shared", "non_sequences"] + } + init_carry["step"] = 0 + return n_steps, sequences, init_carry + + def build_inner_outputs_map(outer_in): + """Map the inner-output variables to their position in the tuple returned by the inner function. + + Inner-outputs are ordered as follow: + - mit-mot-outputs + - mit-sot-outputs + - sit-sot-outputs + - nit-sots (no carry) + - shared-outputs + [+ while-condition] + + """ + inner_outputs_names = ["mit_sot", "sit_sot", "nit_sot", "shared"] + + offset = 0 + inner_output_idx = defaultdict(list) + for name in inner_outputs_names: + num_outputs = len(outer_in[name]) + for i in range(num_outputs): + inner_output_idx[name].append(offset + i) + offset += num_outputs + + return inner_output_idx + + def from_carry_storage(carry, step, input_taps): + """Fetch the inner inputs from the values stored in the carry array. + + `Scan` passes storage arrays as inputs, which are then read from and + updated in the loop body. At each step we need to read from this array + the inputs that will be passed to the inner function. + + This mechanism is necessary because we handle multiple-input taps within + the `scan` instead of letting users manage the memory in the use cases + where this is necessary. + + """ + + def fetch(carry, step, offset): + return carry[step + offset] + + inner_inputs = [] + for taps, carry_element in zip(input_taps, carry): + storage_size = -min(taps) + offsets = [storage_size + tap for tap in taps] + inner_inputs.append( + [fetch(carry_element, step, offset) for offset in offsets] + ) + + return sum(inner_inputs, []) + + def to_carry_storage(inner_outputs, carry, step, input_taps): + """Create the new carry array from the inner output + + `Scan` passes storage arrays as inputs, which are then read from and + updated in the loop body. At each step we need to update this array + with the outputs of the inner function + + """ + new_carry_element = [] + for taps, carry_element, output in zip(input_taps, carry, inner_outputs): + new_carry_element.append( + [carry_element.at[step - tap].set(output) for tap in taps] + ) + + return sum(new_carry_element, []) + + def scan(*outer_inputs): + + outer_in = parse_outer_inputs(outer_inputs) + n_steps, sequences, init_carry = build_jax_scan_inputs(outer_in) + inner_output_idx = build_inner_outputs_map(outer_in) + + def scan_inner_in_args(carry, x): + """Get inner-inputs from the arguments passed to the `jax.lax.scan` body function. + + Inner-inputs are ordered as follows: + - sequences + - mit-mot inputs + - mit-sot inputs + - sit-sot inputs + - shared-inputs + - non-sequences + + """ + current_step = carry["step"] + inner_in_seqs = x + inner_in_mit_sot = from_carry_storage( + carry["mit_sot"], current_step, input_taps["mit_sot"] + ) + inner_in_sit_sot = from_carry_storage( + carry["sit_sot"], current_step, input_taps["sit_sot"] + ) + inner_in_shared = carry.get("shared", []) + inner_in_non_sequences = carry.get("non_sequences", []) - # `inner_scan_inputs` is a list with the following composite form: - # inner_in_seqs - # + sum(inner_in_mit_mot, []) - # + sum(inner_in_mit_sot, []) - # + inner_in_sit_sot - # + inner_in_shared - # + inner_in_non_seqs - inner_in_mit_sot_flatten = [] - for array, index in zip(inner_in_mit_sot, scan_args.mit_sot_in_slices): - inner_in_mit_sot_flatten.extend(array[jnp.array(index)]) - - inner_scan_inputs = sum( + return sum( [ inner_in_seqs, - inner_in_mit_mot, - inner_in_mit_sot_flatten, + inner_in_mit_sot, inner_in_sit_sot, inner_in_shared, - inner_in_non_seqs, + inner_in_non_sequences, ], [], ) - return inner_scan_inputs - - def inner_scan_outs_to_jax_outs( - op, - old_carry, - inner_scan_outs, - ): - ( - inner_in_mit_mot, - inner_in_mit_sot, - inner_in_sit_sot, - inner_in_shared, - inner_in_non_seqs, - ) = old_carry - - def update_mit_sot(mit_sot, new_val): - return jnp.concatenate([mit_sot[1:], new_val[None, ...]], axis=0) - - inner_out_mit_sot = [ - update_mit_sot(mit_sot, new_val) - for mit_sot, new_val in zip(inner_in_mit_sot, inner_scan_outs) - ] - - # This should contain all inner-output taps, non_seqs, and shared - # terms - if not inner_in_sit_sot: - inner_out_sit_sot = [] - else: - inner_out_sit_sot = inner_scan_outs - new_carry = ( - inner_in_mit_mot, - inner_out_mit_sot, - inner_out_sit_sot, - inner_in_shared, - inner_in_non_seqs, - ) + def scan_new_carry(carry, inner_outputs): + """Create a new carry value from the values returned by the inner function (inner-outputs).""" + step = carry["step"] + new_carry = { + "mit_sot": [], + "sit_sot": [], + "shared": [], + "step": step + 1, + "non_sequences": carry["non_sequences"], + } + + if "shared" in inner_output_idx: + shared_inner_outputs = [ + inner_outputs[idx] for idx in inner_output_idx["shared"] + ] + new_carry["shared"] = shared_inner_outputs + + if "mit_sot" in inner_output_idx: + mit_sot_inner_outputs = [ + inner_outputs[idx] for idx in inner_output_idx["mit_sot"] + ] + new_carry["mit_sot"] = to_carry_storage( + mit_sot_inner_outputs, carry["mit_sot"], step, input_taps["mit_sot"] + ) + + if "sit_sot" in inner_output_idx: + sit_sot_inner_outputs = [ + inner_outputs[idx] for idx in inner_output_idx["sit_sot"] + ] + new_carry["sit_sot"] = to_carry_storage( + sit_sot_inner_outputs, carry["sit_sot"], step, input_taps["sit_sot"] + ) return new_carry - def jax_inner_func(carry, x): - inner_args = jax_args_to_inner_scan(op, carry, x) - inner_scan_outs = list(jax_at_inner_func(*inner_args)) - new_carry = inner_scan_outs_to_jax_outs(op, carry, inner_scan_outs) - return new_carry, inner_scan_outs - - _, scan_out = jax.lax.scan(jax_inner_func, init_carry, seqs, length=n_steps) - - # We need to prepend the initial values so that the JAX output will - # match the raw `Scan` `Op` output and, thus, work with a downstream - # `Subtensor` `Op` introduced by the `scan` helper function. - def append_scan_out(scan_in_part, scan_out_part): - return jnp.concatenate([scan_in_part[:-n_steps], scan_out_part], axis=0) - - if scan_args.outer_in_mit_sot: - scan_out_final = [ - append_scan_out(init, out) - for init, out in zip(scan_args.outer_in_mit_sot, scan_out) - ] - elif scan_args.outer_in_sit_sot: - scan_out_final = [ - append_scan_out(init, out) - for init, out in zip(scan_args.outer_in_sit_sot, scan_out) - ] - - if len(scan_out_final) == 1: - scan_out_final = scan_out_final[0] - return scan_out_final + def scan_new_outputs(inner_outputs): + """Create a new outer-output value from the outputs of the inner function. + + Outer-outputs are ordered as follows: + - mit-mot-outputs + - mit-sot-outputs + - sit-sot-outputs + - nit-sots + - shared-outputs + + The shared output corresponds to the last value found in the last + carry value returned by `jax.lax.scan`. It is thus not returned in + the body function. + + """ + outer_outputs = [] + if "mit_sot" in inner_output_idx: + outer_outputs.append( + [inner_outputs[idx] for idx in inner_output_idx["mit_sot"]] + ) + if "sit_sot" in inner_output_idx: + outer_outputs.append( + [inner_outputs[idx] for idx in inner_output_idx["sit_sot"]] + ) + if "nit_sot" in inner_output_idx: + outer_outputs.append( + [inner_outputs[idx] for idx in inner_output_idx["nit_sot"]] + ) + + return tuple(sum(outer_outputs, [])) + + def body_fn(carry, x): + inner_in_args = scan_inner_in_args(carry, x) + inner_outputs = scan_inner_fn(*inner_in_args) + new_carry = scan_new_carry(carry, inner_outputs) + outer_outputs = scan_new_outputs(inner_outputs) + return new_carry, outer_outputs + + last_carry, results = jax.lax.scan( + body_fn, init_carry, sequences, length=n_steps + ) + + shared_output = tuple(last_carry["shared"]) + results = results + shared_output + + if len(results) == 1: + return results[0] + + return results return scan diff --git a/tests/link/jax/test_scan.py b/tests/link/jax/test_scan.py index 158f8bd14d..0992ccd47a 100644 --- a/tests/link/jax/test_scan.py +++ b/tests/link/jax/test_scan.py @@ -3,16 +3,322 @@ from packaging.version import parse as version_parse import aesara.tensor as at +from aesara import function +from aesara.compile.mode import Mode from aesara.configdefaults import config from aesara.graph.fg import FunctionGraph +from aesara.graph.rewriting.db import RewriteDatabaseQuery +from aesara.link.jax.linker import JAXLinker from aesara.scan.basic import scan +from aesara.scan.op import Scan from aesara.tensor.math import gammaln, log +from aesara.tensor.random.utils import RandomStream from aesara.tensor.type import ivector, lscalar, scalar from tests.link.jax.test_basic import compare_jax_and_py jax = pytest.importorskip("jax") +# Disable all optimizations +opts = RewriteDatabaseQuery(include=[None], exclude=["cxx_only", "BlasOpt"]) +jax_no_opts = Mode(JAXLinker(), opts) +py_no_opts = Mode("py", opts) + + +def test_while_cannnot_use_all_outputs(): + """The JAX backend cannot use all the outputs of a while loop. + + Indeed, JAX has fundamental limitations that prevent it from returning + all the intermediate results computed in a `jax.lax.while_loop` loop. + """ + res, updates = scan( + fn=lambda a_tm1: (a_tm1 + 1, until(a_tm1 > 2)), + outputs_info=[{"initial": at.as_tensor(1, dtype=np.int64), "taps": [-1]}], + n_steps=5, + ) + with pytest.raises(NotImplementedError): + function((), res, updates=updates, mode="JAX") + + +def test_while_only_last_output(): + """Compile a `Scan` used as a while loop when only the last computed value + is used. + + """ + res, updates = scan( + fn=lambda a_tm1: (a_tm1 + 1, until(a_tm1 > 2)), + outputs_info=[{"initial": at.as_tensor(1, dtype=np.int64), "taps": [-1]}], + n_steps=5, + ) + res = res[-1] + + jax_fn = function((), res, updates=updates, mode="JAX") + fn = function((), res, updates=updates) + assert np.allclose(fn(), jax_fn()) + + +@pytest.mark.xfail( + reason="Elemwise{add} transforms concrete values into `TracedArray`s" +) +def test_sit_sot(): + a_at = at.scalar("a", dtype="floatX") + + res, updates = scan( + fn=lambda a_tm1: 2 * a_tm1, + outputs_info=[{"initial": a_at, "taps": [-1]}], + n_steps=3, + ) + + fn = function((a_at,), res, updates=updates) + jax_fn = function((a_at,), res, updates=updates, mode=jax_no_opts) + assert np.allclose(fn(1.0), jax_fn(1.0)) + + +def test_sit_sot_opt(): + a_at = at.scalar("a", dtype="floatX") + + res, updates = scan( + fn=lambda a_tm1: 2 * a_tm1, + outputs_info=[{"initial": a_at, "taps": [-1]}], + n_steps=3, + ) + + jax_fn = function((a_at,), res, updates=updates, mode="JAX") + fn = function((a_at,), res, updates=updates) + assert np.allclose(fn(1.0), jax_fn(1.0)) + + +def test_nit_sot_shared(): + res, updates = scan( + fn=lambda: RandomStream(seed=1930, rng_ctor=np.random.RandomState).normal( + 0, 1, name="a" + ), + n_steps=3, + ) + + jax_fn = function((), res, updates=updates, mode="JAX") + res_jax = jax_fn() + fn = function((), res, updates=updates) + res = fn() + + assert res_jax.shape == res.shape + assert not np.all(res_jax == res_jax[0]) + + +@pytest.mark.xfail( + reason="Elemwise{add} transforms concrete values into `TracedArray`s" +) +def test_mit_sot(): + res, updates = scan( + fn=lambda a_tm1: 2 * a_tm1, + outputs_info=[ + {"initial": at.as_tensor([0.0, 1.0], dtype="floatX"), "taps": [-2]} + ], + n_steps=6, + ) + + jax_fn = function((), res, updates=updates, mode=jax_no_opts) + fn = function((), res, updates=updates) + assert np.allclose(fn(), jax_fn()) + + +def test_mit_sot_opt(): + res, updates = scan( + fn=lambda a_tm1: 2 * a_tm1, + outputs_info=[ + {"initial": at.as_tensor([0.0, 1.0], dtype="floatX"), "taps": [-2]} + ], + n_steps=6, + ) + + jax_fn = function((), res, updates=updates, mode="JAX") + fn = function((), res, updates=updates) + assert np.allclose(fn(), jax_fn()) + + +@pytest.mark.xfail( + reason="Elemwise{add} transforms concrete values into `TracedArrays`" +) +def test_mit_sot_2(): + res, updates = scan( + fn=lambda a_tm1, b_tm1: (2 * a_tm1, 2 * b_tm1), + outputs_info=[ + {"initial": at.as_tensor(1.0, dtype="floatX"), "taps": [-1]}, + {"initial": at.as_tensor(0.5, dtype="floatX"), "taps": [-1]}, + ], + n_steps=10, + ) + jax_fn = function((), res, updates=updates, mode=jax_no_opts) + fn = function((), res, updates=updates) + assert np.allclose(fn(), jax_fn()) + + +def test_mit_sot_2_opt(): + res, updates = scan( + fn=lambda a_tm1, b_tm1: (2 * a_tm1, 2 * b_tm1), + outputs_info=[ + {"initial": at.as_tensor(1.0, dtype="floatX"), "taps": [-1]}, + {"initial": at.as_tensor(0.5, dtype="floatX"), "taps": [-1]}, + ], + n_steps=10, + ) + jax_fn = function((), res, updates=updates, mode="JAX") + fn = function((), res, updates=updates) + assert np.allclose(fn(), jax_fn()) + + +@pytest.mark.xfail(reason="Indexing with non-static values in the optimized graph") +def test_sequence_opt(): + a_at = at.dvector("a") + res, updates = scan(fn=lambda a_t: 2 * a_t, sequences=a_at) + jax_fn = function((a_at,), res, updates=updates, mode="JAX") + fn = function((a_at,), res, updates=updates) + assert np.allclose(fn(np.arange(10)), jax_fn(np.arange(10))) + + +@pytest.mark.parametrize("jax_mode", ("JAX", jax_no_opts)) +@pytest.mark.parametrize( + "fn, sequences, outputs_info, non_sequences, n_steps, input_vals, output_vals, op_check", + [ + # sequences + # ( + # lambda a_t: 2 * a_t, + # [at.dvector("a")], + # [{}], + # [], + # None, + # [np.arange(10)], + # None, + # lambda op: op.info.n_seqs > 0, + # ), + # # nit-sot + ( + lambda: at.as_tensor(2.0), + [], + [{}], + [], + 3, + [], + None, + lambda op: op.info.n_nit_sot > 0, + ), + # nit-sot, non_seq + ( + lambda c: at.as_tensor(2.0) * c, + [], + [{}], + [at.dscalar("c")], + 3, + [1.0], + None, + lambda op: op.info.n_nit_sot > 0 and op.info.n_non_seqs > 0, + ), + # sit-sot + # ( + # lambda a_tm1: 2 * a_tm1, + # [], + # [{"initial": at.as_tensor(0.0, dtype="floatX"), "taps": [-1]}], + # [], + # 3, + # [], + # lambda op: op.info.n_sit_sot > 0, + # ), + # # sit-sot, while + # ( + # lambda a_tm1: (a_tm1 + 1, until(a_tm1 > 2)), + # [], + # [{"initial": at.as_tensor(1, dtype=np.int64), "taps": [-1]}], + # [], + # 3, + # [], + # None, + # lambda op: op.info.n_sit_sot > 0, + # ), + # # nit-sot, shared input/output + ( + lambda: RandomStream(seed=1930, rng_ctor=np.random.RandomState).normal( + 0, 1, name="a" + ), + [], + [{}], + [], + 3, + [], + [np.array([-0.4587753, -0.89655604, 2.13323775])], + lambda op: op.info.n_shared_outs > 0, + ), + # mit-sot (that's also a type of sit-sot) + # ( + # lambda a_tm1: 2 * a_tm1, + # [], + # [{"initial": at.as_tensor([0.0, 1.0], dtype="floatX"), "taps": [-2]}], + # [], + # 6, + # [], + # None, + # lambda op: op.info.n_mit_sot > 0, + # ), + # # mit-sot + # ( + # lambda a_tm1, b_tm1: (2 * a_tm1, 2 * b_tm1), + # [], + # [ + # {"initial": at.as_tensor(1.0, dtype="floatX"), "taps": [-1]}, + # {"initial": at.as_tensor(0.3, dtype="floatX"), "taps": [-1]}, + # ], + # [], + # 10, + # [], + # None, + # lambda op: op.info.n_mit_sot > 0, + # ), + ], +) +def test_xit_xot_types( + jax_mode, + fn, + sequences, + outputs_info, + non_sequences, + n_steps, + input_vals, + output_vals, + op_check, +): + """Test basic xit-xot configurations.""" + res, updates = scan( + fn, + sequences=sequences, + outputs_info=outputs_info, + non_sequences=non_sequences, + n_steps=n_steps, + strict=True, + ) + + if not isinstance(res, list): + res = [res] + + # Get rid of any `Subtensor` indexing on the `Scan` outputs + res = [r.owner.inputs[0] if not isinstance(r.owner.op, Scan) else r for r in res] + + scan_op = res[0].owner.op + assert isinstance(scan_op, Scan) + + _ = op_check(scan_op) + + if output_vals is None: + compare_jax_and_py( + ((sequences + non_sequences), res), + input_vals, + updates=updates, + ) + else: + jax_fn = function( + (sequences + non_sequences), res, updates=updates, mode=jax_mode + ) + res_vals = jax_fn(*input_vals) + assert np.allclose(res_vals, output_vals) + @pytest.mark.xfail( version_parse(jax.__version__) >= version_parse("0.2.12"), @@ -144,3 +450,29 @@ def input_step_fn(y_tm1, y_tm3, a): test_input_vals = [np.array(10.0).astype(config.floatX)] compare_jax_and_py(out_fg, test_input_vals) + + +def test_scan_multiple_none_output(): + A = at.dvector("A") + + def power_step(prior_result, x): + return prior_result * x, prior_result * x * x, prior_result * x * x * x + + result, _ = scan( + power_step, + non_sequences=[A], + outputs_info=[at.ones_like(A), None, None], + n_steps=3, + ) + + FunctionGraph([A], result) + test_input_vals = (np.array([1.0, 2.0]),) + + jax_fn = function((A,), result, mode="JAX") + jax_res = jax_fn(*test_input_vals) + + fn = function((A,), result) + res = fn(*test_input_vals) + + for output_jax, output in zip(jax_res, res): + assert np.allclose(jax_res, res) From 2fd81224c5200f769c722c03f5cdce03b5be3145 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?R=C3=A9mi=20Louf?= Date: Tue, 1 Nov 2022 18:37:42 +0100 Subject: [PATCH 3/6] Support `while` loops in the JAX `Scan` dispatcher --- aesara/link/jax/dispatch/scan.py | 311 ++++++++++++++++++++++++++++-- aesara/link/jax/dispatch/shape.py | 2 +- aesara/link/jax/linker.py | 6 +- tests/link/jax/test_scan.py | 26 +-- 4 files changed, 320 insertions(+), 25 deletions(-) diff --git a/aesara/link/jax/dispatch/scan.py b/aesara/link/jax/dispatch/scan.py index 554939d1f2..79ae894128 100644 --- a/aesara/link/jax/dispatch/scan.py +++ b/aesara/link/jax/dispatch/scan.py @@ -2,18 +2,64 @@ from typing import Callable, Dict, List import jax +import jax.numpy as jnp from aesara.link.jax.dispatch.basic import jax_funcify from aesara.scan.op import Scan +from aesara.tensor.shape import Shape_i +from aesara.tensor.subtensor import Subtensor from aesara.tensor.var import TensorVariable +def assert_while_returns_last_output(fgraph, node): + """Check that the clients of the `Scan` outputs are `Subtensor` operators. + + JAX cannot accumulate the intermediate values in a `jax.lax.while` loop, we + thus cannot in general compile a `Scan` operator used as a while loop. + However, when only the last output of the `Scan` computation is used in the + rest of the graph we can transpile `Scan` directly to `jax.lax.while`. + + """ + msg = """JAX cannot accumulate the results inside a `jax.lax.while_loop` loop. + + As a result, Aesara cannot compile the graph you provided to code that can + be run with JAX. In case you only need the value computed at the last iteration + of the while loop, extract this value from the output of `aesara.scan` explicitly + like so: + + >>> res, _ = aesara.scan(...) + >>> value = res[-1] + + So Aesara can compile the graph to code that can be run with JAX. + + """ + + # Count the number of outputs of the outer function. We ignore + # `shared` variables since they are not not accumulated and not + # returned to the user. + op = node.op + num_outer_outputs = ( + op.info.n_mit_mot + op.info.n_mit_sot + op.info.n_sit_sot + op.info.n_nit_sot + ) + for out in node.outputs[:num_outer_outputs]: + for client, _ in fgraph.clients[out]: + if isinstance(client, str): + raise NotImplementedError(msg) + elif isinstance(client.op, Subtensor): + idx_list = client.op.idx_list + if isinstance(idx_list[0], slice): + raise NotImplementedError(msg) + elif not isinstance(client.op, Shape_i): + raise NotImplementedError(msg) + + @jax_funcify.register(Scan) def jax_funcify_Scan(op, node, **kwargs): scan_inner_fn = jax_funcify(op.fgraph) input_taps = { "mit_sot": op.info.mit_sot_in_slices, "sit_sot": op.info.sit_sot_in_slices, + "nit_sot": op.info.sit_sot_in_slices, } # Outer-inputs are the inputs to the `Scan` apply node, built from the @@ -36,13 +82,255 @@ def parse_outer_inputs(outer_inputs): return outer_in if op.info.as_while: - raise NotImplementedError("While loops are not supported in the JAX backend.") + global_fgraph = kwargs.get("global_fgraph", None) + assert_while_returns_last_output(global_fgraph, node) + return make_jax_while_fn(scan_inner_fn, parse_outer_inputs, input_taps) else: - return make_jax_scan_fn( - scan_inner_fn, - parse_outer_inputs, - input_taps, - ) + return make_jax_scan_fn(scan_inner_fn, parse_outer_inputs, input_taps) + + +def make_jax_while_fn( + scan_inner_fn: Callable, + parse_outer_inputs: Callable[[TensorVariable], Dict[str, List[TensorVariable]]], + input_taps: Dict, +): + """Create a `jax.lax.while_loop` function to perform `Scan` computations when it + is used as while loop. + + `jax.lax.while_loop` iterates by passing a value `carry` to a `body_fun` that + must return a value of the same type (Pytree structure, shape and dtype of + the leaves). Before calling `body_fn`, it calls `cond_fn` which takes the + current value and returns a boolean that indicates whether to keep iterating + or not. + + The JAX `while_loop` needs to perform the following operations: + + 1. Extract the inner-inputs; + 2. Build the initial carry value; + 3. Inside the loop: + 1. `carry` -> inner-inputs; + 2. inner-outputs -> `carry` + 4. Post-process the `carry` storage and return outputs + """ + + def build_while_carry(outer_in): + """Build the inputs to `jax.lax.scan` from the outer-inputs.""" + init_carry = { + "mit_sot": [], + "mit_sot_storage": outer_in["mit_sot"], + "sit_sot": [], + "sit_sot_storage": outer_in["sit_sot"], + "shared": outer_in["shared"], + "sequences": outer_in["sequences"], + "non_sequences": outer_in["non_sequences"], + } + init_carry["step"] = 0 + init_carry["do_stop"] = False + return init_carry + + def build_inner_outputs_map(outer_in): + """Map the inner-output variables to their position in the tuple returned by the inner function. + + TODO: Copied from the scan builder + + Inner-outputs are ordered as follow: + - mit-mot-outputs + - mit-sot-outputs + - sit-sot-outputs + - nit-sots (no carry) + - shared-outputs + [+ while-condition] + + """ + inner_outputs_names = ["mit_sot", "sit_sot", "nit_sot", "shared"] + + offset = 0 + inner_output_idx = defaultdict(list) + for name in inner_outputs_names: + num_outputs = len(outer_in[name]) + for i in range(num_outputs): + inner_output_idx[name].append(offset + i) + offset += num_outputs + + return inner_output_idx + + def from_carry_storage(carry, step, input_taps): + """Fetch the inner inputs from the values stored in the carry array. + + `Scan` passes storage arrays as inputs, which are then read from and + updated in the loop body. At each step we need to read from this array + the inputs that will be passed to the inner function. + + This mechanism is necessary because we handle multiple-input taps within + the `scan` instead of letting users manage the memory in the use cases + where this is necessary. + + TODO: Copied from the scan builder + + """ + + def fetch(carry, step, offset): + return carry[step + offset] + + inner_inputs = [] + for taps, carry_element in zip(input_taps, carry): + storage_size = -min(taps) + offsets = [storage_size + tap for tap in taps] + inner_inputs.append( + [fetch(carry_element, step, offset) for offset in offsets] + ) + + return sum(inner_inputs, []) + + def to_carry_storage(inner_outputs, carry, step, input_taps): + """Create the new carry array from the inner output + + `Scan` passes storage arrays as inputs, which are then read from and + updated in the loop body. At each step we need to update this array + with the outputs of the inner function + + TODO: Copied from the scan builder + + """ + new_carry_element = [] + for taps, carry_element, output in zip(input_taps, carry, inner_outputs): + new_carry_element.append( + [carry_element.at[step - tap].set(output) for tap in taps] + ) + + return sum(new_carry_element, []) + + def while_loop(*outer_inputs): + + outer_in = parse_outer_inputs(outer_inputs) + init_carry = build_while_carry(outer_in) + inner_output_idx = build_inner_outputs_map(outer_in) + + def inner_inputs_from_carry(carry): + """Get inner-inputs from the arguments passed to the `jax.lax.while_loop` body function. + + Inner-inputs are ordered as follows: + - sequences + - mit-mot inputs + - mit-sot inputs + - sit-sot inputs + - shared-inputs + - non-sequences + + """ + current_step = carry["step"] + + inner_in_mit_sot = from_carry_storage( + carry["mit_sot_storage"], current_step, input_taps["mit_sot"] + ) + inner_in_sit_sot = from_carry_storage( + carry["sit_sot_storage"], current_step, input_taps["sit_sot"] + ) + inner_in_shared = carry.get("shared", []) + inner_in_non_sequences = carry.get("non_sequences", []) + + return sum( + [ + inner_in_mit_sot, + inner_in_sit_sot, + inner_in_shared, + inner_in_non_sequences, + ], + [], + ) + + def carry_from_inner_outputs(carry, inner_outputs): + step = carry["step"] + new_carry = { + "mit_sot": [], + "sit_sot": [], + "sit_sot_storage": [], + "nit_sot": [], + "mit_sot_storage": [], + "shared": [], + "step": step + 1, + "sequences": carry["sequences"], + "non_sequences": carry["non_sequences"], + "do_stop": inner_outputs[-1], + } + + if "shared" in inner_output_idx: + shared_inner_outputs = [ + inner_outputs[idx] for idx in inner_output_idx["shared"] + ] + new_carry["shared"] = shared_inner_outputs + + if "mit_sot" in inner_output_idx: + mit_sot_inner_outputs = [ + inner_outputs[idx] for idx in inner_output_idx["mit_sot"] + ] + new_carry["mit_sot"] = mit_sot_inner_outputs + new_carry["mit_sot_storage"] = to_carry_storage( + mit_sot_inner_outputs, + carry["mit_sot_storage"], + step, + input_taps["mit_sot"], + ) + + if "sit_sot" in inner_output_idx: + sit_sot_inner_outputs = [ + inner_outputs[idx] for idx in inner_output_idx["sit_sot"] + ] + new_carry["sit_sot"] = sit_sot_inner_outputs + new_carry["sit_sot_storage"] = to_carry_storage( + sit_sot_inner_outputs, + carry["sit_sot_storage"], + step, + input_taps["sit_sot"], + ) + + if "nit_sot" in inner_output_idx: + nit_sot_inner_outputs = [ + inner_outputs[idx] for idx in inner_output_idx["nit_sot"] + ] + new_carry["nit_sot"] = nit_sot_inner_outputs + + return new_carry + + def cond_fn(carry): + # The inner-function of `Scan` returns a boolean as the last + # value. This needs to be included in `carry`. + # TODO: Will it return `False` if the number of steps is exceeded? + return ~carry["do_stop"] + + def body_fn(carry): + inner_inputs = inner_inputs_from_carry(carry) + inner_outputs = scan_inner_fn(*inner_inputs) + new_carry = carry_from_inner_outputs(carry, inner_outputs) + return new_carry + + # The `Scan` implementation in the C backend will execute the + # function once before checking the termination condition, while + # `jax.lax.while_loop` checks the condition first. We thus need to call + # `body_fn` once before calling `jax.lax.while_loop`. This allows us, + # along with `n_steps`, to build the storage array for the `nit-sot`s + # since there is no way to know their shape and dtype before executing + # the function. + inner_inputs = inner_inputs_from_carry(init_carry) + inner_outputs = scan_inner_fn(*inner_inputs) + carry = carry_from_inner_outputs(init_carry, inner_outputs) + carry = jax.lax.while_loop(cond_fn, body_fn, carry) + + # Post-process the storage arrays + # We make sure that the outputs are not scalars in case an array + # is expected downstream since `Scan` is supposed to always return arrays + carry["sit_sot"] = [jnp.atleast_1d(element) for element in carry["sit_sot"]] + carry["mit_sot"] = [jnp.atleast_1d(element) for element in carry["mit_sot"]] + carry["nit_not"] = [jnp.atleast_1d(element) for element in carry["nit_sot"]] + + outer_outputs = ["mit_sot", "sit_sot", "nit_sot", "shared"] + results = sum([carry[output] for output in outer_outputs], []) + if len(results) == 1: + return results[0] + else: + return results + + return while_loop def make_jax_scan_fn( @@ -58,7 +346,8 @@ def make_jax_scan_fn( stacked to the previous outputs. We use this to our advantage to build `Scan` outputs without having to post-process the storage arrays. - The JAX scan function needs to perform the following operations: + The JAX `scan` function needs to perform the following operations: + 1. Extract the inner-inputs; 2. Build the initial `carry` and `sequence` values; 3. Inside the loop: @@ -265,11 +554,11 @@ def body_fn(carry, x): ) shared_output = tuple(last_carry["shared"]) - results = results + shared_output + outer_outputs = results + shared_output - if len(results) == 1: - return results[0] + if len(outer_outputs) == 1: + return outer_outputs[0] - return results + return outer_outputs return scan diff --git a/aesara/link/jax/dispatch/shape.py b/aesara/link/jax/dispatch/shape.py index 6fd7a74fb1..f982269445 100644 --- a/aesara/link/jax/dispatch/shape.py +++ b/aesara/link/jax/dispatch/shape.py @@ -84,7 +84,7 @@ def shape(x): @jax_funcify.register(Shape_i) -def jax_funcify_Shape_i(op, **kwargs): +def jax_funcify_Shape_i(op, node, **kwargs): i = op.i def shape_i(x): diff --git a/aesara/link/jax/linker.py b/aesara/link/jax/linker.py index 49ef83b293..51f77412c6 100644 --- a/aesara/link/jax/linker.py +++ b/aesara/link/jax/linker.py @@ -50,7 +50,11 @@ def fgraph_convert(self, fgraph, input_storage, storage_map, **kwargs): ) return jax_funcify( - fgraph, input_storage=input_storage, storage_map=storage_map, **kwargs + fgraph, + input_storage=input_storage, + storage_map=storage_map, + global_fgraph=fgraph, + **kwargs, ) def jit_compile(self, fn): diff --git a/tests/link/jax/test_scan.py b/tests/link/jax/test_scan.py index 0992ccd47a..617e651762 100644 --- a/tests/link/jax/test_scan.py +++ b/tests/link/jax/test_scan.py @@ -11,6 +11,7 @@ from aesara.link.jax.linker import JAXLinker from aesara.scan.basic import scan from aesara.scan.op import Scan +from aesara.scan.utils import until from aesara.tensor.math import gammaln, log from aesara.tensor.random.utils import RandomStream from aesara.tensor.type import ivector, lscalar, scalar @@ -26,10 +27,11 @@ def test_while_cannnot_use_all_outputs(): - """The JAX backend cannot use all the outputs of a while loop. + """The JAX backend cannot return all the outputs of a while loop. Indeed, JAX has fundamental limitations that prevent it from returning all the intermediate results computed in a `jax.lax.while_loop` loop. + """ res, updates = scan( fn=lambda a_tm1: (a_tm1 + 1, until(a_tm1 > 2)), @@ -202,6 +204,16 @@ def test_sequence_opt(): None, lambda op: op.info.n_nit_sot > 0, ), + # ( + # lambda: at.as_tensor(2.0), + # [], + # [{}], + # [], + # 3, + # [], + # None, + # lambda op: op.info.n_nit_sot > 0, + # ), # nit-sot, non_seq ( lambda c: at.as_tensor(2.0) * c, @@ -221,16 +233,6 @@ def test_sequence_opt(): # [], # 3, # [], - # lambda op: op.info.n_sit_sot > 0, - # ), - # # sit-sot, while - # ( - # lambda a_tm1: (a_tm1 + 1, until(a_tm1 > 2)), - # [], - # [{"initial": at.as_tensor(1, dtype=np.int64), "taps": [-1]}], - # [], - # 3, - # [], # None, # lambda op: op.info.n_sit_sot > 0, # ), @@ -258,7 +260,7 @@ def test_sequence_opt(): # None, # lambda op: op.info.n_mit_sot > 0, # ), - # # mit-sot + # mit-sot # ( # lambda a_tm1, b_tm1: (2 * a_tm1, 2 * b_tm1), # [], From 0932c8e85bf419bae72a06dbcfdb516a9e7f9d2a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?R=C3=A9mi=20Louf?= Date: Thu, 8 Dec 2022 14:28:24 +0100 Subject: [PATCH 4/6] Support mit-mots in the JAX backend --- aesara/link/jax/dispatch/scan.py | 25 +++++++++++++++---- tests/link/jax/test_scan.py | 42 +++++++++++++++++++++++++++++--- 2 files changed, 59 insertions(+), 8 deletions(-) diff --git a/aesara/link/jax/dispatch/scan.py b/aesara/link/jax/dispatch/scan.py index 79ae894128..2f93b349ef 100644 --- a/aesara/link/jax/dispatch/scan.py +++ b/aesara/link/jax/dispatch/scan.py @@ -57,6 +57,7 @@ def assert_while_returns_last_output(fgraph, node): def jax_funcify_Scan(op, node, **kwargs): scan_inner_fn = jax_funcify(op.fgraph) input_taps = { + "mit_mot": op.info.mit_mot_in_slices, "mit_sot": op.info.mit_sot_in_slices, "sit_sot": op.info.sit_sot_in_slices, "nit_sot": op.info.sit_sot_in_slices, @@ -76,9 +77,6 @@ def parse_outer_inputs(outer_inputs): "shared": list(op.outer_shared(outer_inputs)), "non_sequences": list(op.outer_non_seqs(outer_inputs)), } - if len(outer_in["mit_mot"]) > 0: - raise NotImplementedError("mit-mot not supported") - return outer_in if op.info.as_while: @@ -364,7 +362,7 @@ def build_jax_scan_inputs(outer_in: Dict): sequences = outer_in["sequences"] init_carry = { name: outer_in[name] - for name in ["mit_sot", "sit_sot", "shared", "non_sequences"] + for name in ["mit_mot", "mit_sot", "sit_sot", "shared", "non_sequences"] } init_carry["step"] = 0 return n_steps, sequences, init_carry @@ -381,7 +379,7 @@ def build_inner_outputs_map(outer_in): [+ while-condition] """ - inner_outputs_names = ["mit_sot", "sit_sot", "nit_sot", "shared"] + inner_outputs_names = ["mit_mot", "mit_sot", "sit_sot", "nit_sot", "shared"] offset = 0 inner_output_idx = defaultdict(list) @@ -456,6 +454,9 @@ def scan_inner_in_args(carry, x): current_step = carry["step"] inner_in_seqs = x + inner_in_mit_mot = from_carry_storage( + carry["mit_mot"], current_step, input_taps["mit_mot"] + ) inner_in_mit_sot = from_carry_storage( carry["mit_sot"], current_step, input_taps["mit_sot"] ) @@ -468,6 +469,7 @@ def scan_inner_in_args(carry, x): return sum( [ inner_in_seqs, + inner_in_mit_mot, inner_in_mit_sot, inner_in_sit_sot, inner_in_shared, @@ -480,6 +482,7 @@ def scan_new_carry(carry, inner_outputs): """Create a new carry value from the values returned by the inner function (inner-outputs).""" step = carry["step"] new_carry = { + "mit_mot": [], "mit_sot": [], "sit_sot": [], "shared": [], @@ -493,6 +496,14 @@ def scan_new_carry(carry, inner_outputs): ] new_carry["shared"] = shared_inner_outputs + if "mit_mot" in inner_output_idx: + mit_mot_inner_outputs = [ + inner_outputs[idx] for idx in inner_output_idx["mit_mot"] + ] + new_carry["mit_mot"] = to_carry_storage( + mit_mot_inner_outputs, carry["mit_mot"], step, input_taps["mit_mot"] + ) + if "mit_sot" in inner_output_idx: mit_sot_inner_outputs = [ inner_outputs[idx] for idx in inner_output_idx["mit_sot"] @@ -527,6 +538,10 @@ def scan_new_outputs(inner_outputs): """ outer_outputs = [] + if "mit_mot" in inner_output_idx: + outer_outputs.append( + [inner_outputs[idx] for idx in inner_output_idx["mit_mot"]] + ) if "mit_sot" in inner_output_idx: outer_outputs.append( [inner_outputs[idx] for idx in inner_output_idx["mit_sot"]] diff --git a/tests/link/jax/test_scan.py b/tests/link/jax/test_scan.py index 617e651762..585d86fae1 100644 --- a/tests/link/jax/test_scan.py +++ b/tests/link/jax/test_scan.py @@ -3,7 +3,7 @@ from packaging.version import parse as version_parse import aesara.tensor as at -from aesara import function +from aesara import function, grad from aesara.compile.mode import Mode from aesara.configdefaults import config from aesara.graph.fg import FunctionGraph @@ -27,11 +27,10 @@ def test_while_cannnot_use_all_outputs(): - """The JAX backend cannot return all the outputs of a while loop. + """The JAX backend cannot use all the outputs of a while loop. Indeed, JAX has fundamental limitations that prevent it from returning all the intermediate results computed in a `jax.lax.while_loop` loop. - """ res, updates = scan( fn=lambda a_tm1: (a_tm1 + 1, until(a_tm1 > 2)), @@ -233,6 +232,16 @@ def test_sequence_opt(): # [], # 3, # [], + # lambda op: op.info.n_sit_sot > 0, + # ), + # # sit-sot, while + # ( + # lambda a_tm1: (a_tm1 + 1, until(a_tm1 > 2)), + # [], + # [{"initial": at.as_tensor(1, dtype=np.int64), "taps": [-1]}], + # [], + # 3, + # [], # None, # lambda op: op.info.n_sit_sot > 0, # ), @@ -478,3 +487,30 @@ def power_step(prior_result, x): for output_jax, output in zip(jax_res, res): assert np.allclose(jax_res, res) + + +@pytest.mark.xfail(reason="Fails for reasons unrelated to `Scan`") +def test_mitmots_basic(): + + init_x = at.dvector() + seq = at.dvector() + + def inner_fct(seq, state_old, state_current): + return state_old * 2 + state_current + seq + + out, _ = scan( + inner_fct, sequences=seq, outputs_info={"initial": init_x, "taps": [-2, -1]} + ) + + g_outs = grad(out.sum(), [seq, init_x]) + + out_fg = FunctionGraph([seq, init_x], g_outs) + + seq_val = np.arange(3) + init_x_val = np.r_[-2, -1] + (seq_val, init_x_val) + + fn = function(out_fg.inputs, out_fg.outputs) + jax_fn = function(out_fg.inputs, out_fg.outputs, mode="JAX") + print(fn(seq_val, init_x_val)) + print(jax_fn(seq_val, init_x_val)) From 1d2cc072543576fe47d1526ea53d492792a8a70e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?R=C3=A9mi=20Louf?= Date: Fri, 16 Dec 2022 10:37:09 +0100 Subject: [PATCH 5/6] Pass through `ScalarFromTensor` when input is scalar --- aesara/link/jax/dispatch/tensor_basic.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/aesara/link/jax/dispatch/tensor_basic.py b/aesara/link/jax/dispatch/tensor_basic.py index feecf3922a..6f57f13cae 100644 --- a/aesara/link/jax/dispatch/tensor_basic.py +++ b/aesara/link/jax/dispatch/tensor_basic.py @@ -136,6 +136,8 @@ def tensor_from_scalar(x): @jax_funcify.register(ScalarFromTensor) def jax_funcify_ScalarFromTensor(op, **kwargs): def scalar_from_tensor(x): + if isinstance(x, (float, int)): + return x return jnp.array(x).flatten()[0] return scalar_from_tensor From b09a40e330b2c1e13dba127dcfed3c59b1080831 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?R=C3=A9mi=20Louf?= Date: Tue, 21 Feb 2023 18:42:58 +0100 Subject: [PATCH 6/6] Preprend the initial state to the result of the `scan` loop --- aesara/link/jax/dispatch/scan.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/aesara/link/jax/dispatch/scan.py b/aesara/link/jax/dispatch/scan.py index 2f93b349ef..7c3f377736 100644 --- a/aesara/link/jax/dispatch/scan.py +++ b/aesara/link/jax/dispatch/scan.py @@ -568,9 +568,26 @@ def body_fn(carry, x): body_fn, init_carry, sequences, length=n_steps ) + # We need to preprend the initial values so the output matches + # the raw `Scan` output. + if len(outer_in["mit_sot"]) > 0: + results = tuple([ + jnp.concatenate([init[:-n_steps], out], axis=0) + for init, out in zip(outer_in["mit_sot"], results) + ]) + # TODO: HERE IS THE REASON WHY test_scan_multiple_none_output + elif len(outer_in["sit_sot"]) > 0: + results = tuple([ + jnp.concatenate([init[:-n_steps], out], axis=0) + for init, out in zip(outer_in["sit_sot"], results) + ]) + + breakpoint() + shared_output = tuple(last_carry["shared"]) outer_outputs = results + shared_output + if len(outer_outputs) == 1: return outer_outputs[0]