jax-ml / jax-ml/oryx

Confusion about "append" mode in `sow`

Open
#78 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
328
Forks
13
Avg merge
49m
Merged PRs (30d)
2

Description

I think I am confused on how to use the "append" mode in `sow`, I would expect the following code

``` python
from oryx.core import sow, reap

def f(x):
x = sow(x + 1.0, tag="tag", name="x", mode="append")
x = sow(x + 1.0, tag="tag", name="x", mode="append")
return x

print(reap(f, tag="tag")(1))
```

To output something similar to `{'x': [2, 3]}`, as the documentation says: "Another option is 'append', in which all sows of the same name will be appended into a growing array.". However, I get the same error as in strict mode, which is

```
Traceback (most recent call last):
File "/home/anton/flarenet/oryx_example_append.py", line 11, in
print(reap(f, tag="tag")(1))
^^^^^^^^^^^^^^^^^^^^^
File "/home/anton/flarenet/.venv/lib/python3.12/site-packages/oryx/core/interpreters/harvest.py", line 868, in wrapped
return call_and_reap(
^^^^^^^^^^^^^^
File "/home/anton/flarenet/.venv/lib/python3.12/site-packages/oryx/core/interpreters/harvest.py", line 801, in wrapped
out, reaps, preds = _call_and_reap(
^^^^^^^^^^^^^^^
File "/home/anton/flarenet/.venv/lib/python3.12/site-packages/oryx/core/interpreters/harvest.py", line 838, in wrapped
out_flat, reaps, preds = flat_fun.call_wrapped(flat_args)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/anton/flarenet/.venv/lib/python3.12/site-packages/jax/_src/linear_util.py", line 192, in call_wrapped
ans = self.f(*args, **dict(self.params, **kwargs))
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/anton/flarenet/oryx_example_append.py", line 7, in f
x = sow(x + 1.0, tag="tag", name="x", mode="append")
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/anton/flarenet/.venv/lib/python3.12/site-packages/oryx/core/interpreters/harvest.py", line 253, in sow
return _sow(value, tag=tag, name=name, mode=mode, key=key)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/anton/flarenet/.venv/lib/python3.12/site-packages/oryx/core/interpreters/harvest.py", line 301, in _sow
out_flat = sow_p.bind(*flat_args, name=name, tag=tag, mode=mode, tree=in_tree)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/anton/flarenet/.venv/lib/python3.12/site-packages/jax/_src/core.py", line 416, in bind
return self.bind_with_trace(find_top_trace(args), args, params)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/anton/flarenet/.venv/lib/python3.12/site-packages/jax/_src/core.py", line 420, in bind_with_trace
out = trace.process_primitive(self, map(trace.full_raise, args), params)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/anton/flarenet/.venv/lib/python3.12/site-packages/oryx/core/interpreters/harvest.py", line 396, in process_primitive
return self.default_process_primitive(primitive, tracers, params)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/anton/flarenet/.venv/lib/python3.12/site-packages/oryx/core/interpreters/harvest.py", line 404, in default_process_primitive
outvals = context.process_sow(*vals, **params)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/anton/flarenet/.venv/lib/python3.12/site-packages/oryx/core/interpreters/harvest.py", line 508, in process_sow
return self.handle_sow(*values, name=name, tag=tag, tree=tree, mode=mode)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/home/anton/flarenet/.venv/lib/python3.12/site-packages/oryx/core/interpreters/harvest.py", line 567, in handle_sow
raise ValueError(f'Variable has already been reaped: {name}')
ValueError: Variable has already been reaped: x
```

I am using `python 3.12.3` with `jax==0.4.29` and `jaxlib==0.4.29`. The same code with `mode` set to "clobber" works as described in the documentation.
I would like to have a way to use the same name and tag multiple times, but still be able to `reap` or `plant` in/from (specific) parts of the `sow`ed places, since I do not want to force users to call functions with `sow` in them with `nest`.
What is the recommended way to do this and could you perhaps provide a minimal working example using the append `mode`?

Please let me know if you would like any further information.
Thank you very much!

Contributor guide

Open the contributing guide

Research direction

Reproduce the example with the versions shown, then inspect oryx/core/interpreters/harvest.py, especially sow and handle_sow. Compare the append-mode behavior with the documentation and determine whether the implementation or the example is wrong. Done means the expected repeated-name behavior is clarified with a minimal working example, or a tested fix is identified.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
backend-api-design
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.