mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-25 17:26:08 +00:00
Code Issues Packages Projects Releases Wiki Activity
691 Commits 508 Branches 263 Tags
Commit Graph

12 Commits

Author SHA1 Message Date
Peter Hawkins
2b383bdbd9 Increase cache size to 4096. 2019-01-15 10:32:58 -05:00
Peter Hawkins
3266bb3122 Change linear_util.memoize to use an LRU cache. 
Add util.OrderedDict that retrofits a move_to_end method onto Python 2 OrderedDicts.
 2019-01-14 21:48:28 -05:00
Peter Hawkins
05b1049e49 Change util.memoize to be an LRU cache with a default size of 64 entries. 
The goal is to limit peak memory when running a large number of computations, e.g., the test suite.
 2019-01-14 20:11:08 -05:00
Peter Hawkins
8127392a18 Use get() rather than a try-catch block in memoized function lookup. 
Currently backtraces often look like this:
```
---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
~/p/jax/jax/util.py in memoized_fun(*args, **kwargs)
    133     try:
--> 134       return cache[key]
    135     except KeyError:

KeyError: ((lu, ShapedArray(int32[2,2])), ())

During handling of the above exception, another exception occurred:

KeyError                                  Traceback (most recent call last)
~/p/jax/jax/util.py in memoized_fun(*args, **kwargs)
    133     try:
--> 134       return cache[key]
    135     except KeyError:

KeyError: ((lu, xla_client.Shape(_dtype=dtype('int32'), _dimensions=(2, 2), _is_tuple=False, _minor_to_major=None)), ())

During handling of the above exception, another exception occurred:

NotImplementedError                       Traceback (most recent call last)
<ipython-input-26-d6c00d50e3c9> in <module>
```

The "during handling of the above exception..." message is mostly a distraction for the user that occurs because we perform the memoized function evaluation inside a `catch` block. By performing the function evaluation outside the catch block, we can get better backtraces without the distraction of the KeyError exception.

```
 2018-12-21 13:32:56 -05:00
Alex Wiltschko
763f860135 Automated detection of unimplemented functions 2018-12-11 11:52:31 -05:00
Dougal Maclaurin
30124b6da1 Added jit transformations to generated functions. Fixed bug in comparing numpy arrays for equality. 2018-12-08 00:03:34 -05:00
Dougal Maclaurin
307d195577 Wrapped static args to jit to be hashed on id. This is conservative but simple and cheap. 2018-12-05 15:55:01 -05:00
Dougal Maclaurin
ca2634ea5d source sync 
PiperOrigin-RevId: 222923229
 2018-11-27 16:51:22 -08:00
Peter Hawkins
5e60639bc5 source sync 
PiperOrigin-RevId: 222452709
 2018-11-21 20:22:54 -08:00
Peter Hawkins
e180f08113 source sync 
PiperOrigin-RevId: 222451919
 2018-11-21 20:22:51 -08:00
Peter Hawkins
fe4edf2839 source sync 
PiperOrigin-RevId: 222449830
 2018-11-21 20:22:49 -08:00
Matthew Johnson
a30e858e59 populating source tree 2018-11-17 18:03:33 -08:00